DataTable update to add state props to manage expanded rows externally

packages/comet-extras/src/components/data-table/data-table.stories.tsx

@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useState } from 'react';
 import { Meta } from '@storybook/react-vite';
 import DataTable, { DataTableProps } from './data-table';
 import { createColumnHelper } from '@tanstack/react-table';
@@ -173,3 +173,37 @@ export const WithExpandableRows = {
   },
   render: (args: DataTableProps) => <DataTable {...args} />,
 };
+
+export const WithExternalExpandedRowState = {
+  args: {
+    id: 'table-parent-state',
+    columns: cols,
+    data,
+    striped: false,
+    sortable: true,
+    sortDir: 'asc',
+    sortCol: 'lastName',
+    pageable: true,
+    pageIndex: 0,
+    pageSize: 3,
+    expandable: true,
+    getChildRows: (row: Person) => row.children,
+    className: 'width-full',
+  },
+  parameters: {
+    docs: {
+      source: {
+        type: 'code',
+      },
+    },
+  },
+  render: (args: DataTableProps) => {
+    const [expandedRows, setExpandedRows] = useState<Record<string, boolean> | true>({});
+    return (
+      <div>
+        <p>External state - expanded rows: [{Object.keys(expandedRows).join()}]</p>
+        <DataTable {...args} parentExpanded={expandedRows} setParentExpanded={setExpandedRows} />
+      </div>
+    );
+  },
+};

packages/comet-extras/src/components/data-table/data-table.test.tsx

@@ -1,5 +1,6 @@
-import React from 'react';
-import { act, render } from '@testing-library/react';
+import React, { useState } from 'react';
+import { act, render, screen } from '@testing-library/react';
+import { userEvent } from '@testing-library/user-event';
 import DataTable from './data-table';
 import { createColumnHelper } from '@tanstack/react-table';
 
@@ -355,6 +356,140 @@ describe('DataTable', () => {
     expect(pageButtons.length).toBeGreaterThan(0); // Should have page buttons
   });
 
+  test('should display all rows initially expanded', () => {
+    const dataWithChildren: Person[] = [
+      {
+        firstName: 'John',
+        lastName: 'Doe',
+        children: [
+          {
+            firstName: 'Johnny',
+            lastName: 'Doe Jr',
+          },
+        ],
+      },
+      {
+        firstName: 'Jane',
+        lastName: 'Smith',
+        children: [
+          {
+            firstName: 'Janie',
+            lastName: 'Smith Jr',
+          },
+        ],
+      },
+      {
+        firstName: 'Bob',
+        lastName: 'Wilson',
+        children: [
+          {
+            firstName: 'Bobby',
+            lastName: 'Wilson Jr',
+          },
+        ],
+      },
+    ];
+
+    const { baseElement } = render(
+      <DataTable
+        id="table-all-expanded"
+        columns={cols}
+        data={dataWithChildren}
+        expandable={true}
+        initialExpanded={true}
+        getChildRows={(row: Person) => row.children}
+      ></DataTable>,
+    );
+
+    // Should show 3 parent rows, each with 1 child row
+    const allRows = baseElement.querySelectorAll('tbody tr');
+    expect(allRows).toHaveLength(6); // 3 parents rows + 1 child row each
+
+    // Verify we have 3 child rows
+    const childRows = baseElement.querySelectorAll('.child-row');
+    expect(childRows).toHaveLength(3); // 1 child row each
+  });
+
+  test('external expanded row state', async () => {
+    const dataWithChildren: Person[] = [
+      {
+        firstName: 'John',
+        lastName: 'Doe',
+        children: [
+          {
+            firstName: 'Johnny',
+            lastName: 'Doe Jr',
+          },
+        ],
+      },
+      {
+        firstName: 'Jane',
+        lastName: 'Smith',
+        children: [
+          {
+            firstName: 'Janie',
+            lastName: 'Smith Jr',
+          },
+        ],
+      },
+      {
+        firstName: 'Bob',
+        lastName: 'Wilson',
+        children: [
+          {
+            firstName: 'Bobby',
+            lastName: 'Wilson Jr',
+          },
+        ],
+      },
+    ];
+
+    const Parent = () => {
+      const [expandedRows, setExpandedRows] = useState<Record<string, boolean> | true>({});
+      return (
+        <>
+          <p id="expanded-rows">Expanded rows: [{Object.keys(expandedRows).join()}]</p>
+          <DataTable
+            id="table-all-expanded"
+            columns={cols}
+            data={dataWithChildren}
+            expandable={true}
+            parentExpanded={expandedRows}
+            setParentExpanded={setExpandedRows}
+            getChildRows={(row: Person) => row.children}
+          ></DataTable>
+        </>
+      );
+    };
+
+    const { baseElement } = render(<Parent />);
+
+    // Should show 3 parent rows
+    const allRows = baseElement.querySelectorAll('tbody tr');
+    expect(allRows).toHaveLength(3); // 3 parents rows
+
+    const expandButtons = screen.getAllByRole('button', { name: 'Expand row' });
+    expect(expandButtons).toHaveLength(3); // 3 expandable rows
+
+    await userEvent.click(expandButtons[0]);
+    const parentsPlusOneChild = baseElement.querySelectorAll('tbody tr');
+    expect(parentsPlusOneChild).toHaveLength(4); // 3 parents rows + 1 child
+    // test the external expanded row state is updated correctly
+    let expandedRows = screen.getByText('Expanded rows: [0]');
+    expect(expandedRows).toBeInTheDocument();
+
+    await userEvent.click(expandButtons[2]);
+    const parentsPlusTwoChildren = baseElement.querySelectorAll('tbody tr');
+    expect(parentsPlusTwoChildren).toHaveLength(5); // 3 parents rows + 2 children
+    // test the external expanded row state is updated correctly
+    expandedRows = screen.getByText('Expanded rows: [0,2]');
+    expect(expandedRows).toBeInTheDocument();
+
+    // Verify we have 2 child rows
+    const childRows = baseElement.querySelectorAll('.child-row');
+    expect(childRows).toHaveLength(2); // rows 1 and 3 expanded each with 1 child
+  });
+
   test('should not count child rows against pagination when expanded', async () => {
     const dataWithChildren: Person[] = [
       {

packages/comet-extras/src/components/data-table/data-table.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect } from 'react';
+import React, { Dispatch, SetStateAction, useEffect } from 'react';
 import {
   SortingState,
   flexRender,
@@ -63,9 +63,20 @@ export interface DataTableProps<T = any> {
    */
   getChildRows?: (row: T) => T[] | undefined;
   /**
-   * Initial expanded state for rows (object with row IDs as keys and boolean values)
+   * An optional state value from parent representing the expanded rows.
+   * Value: object with row IDs as keys and boolean values, or true to expand all.
    */
-  initialExpanded?: Record<string, boolean>;
+  parentExpanded?: Record<string, boolean> | true;
+  /**
+   * An optional state setter function from parent for the expanded rows state.
+   */
+  setParentExpanded?: Dispatch<SetStateAction<Record<string, boolean> | true>>;
+  /**
+   * Initial expanded state for rows (object with row IDs as keys and boolean values
+   * or true to expand all).
+   * Only used if parent state value and setter function are not provided.
+   */
+  initialExpanded?: Record<string, boolean> | true;
   /**
    * Additional class names for the table
    */
@@ -91,6 +102,8 @@ export const DataTable = ({
   pageSize = 10,
   expandable = false,
   getChildRows,
+  parentExpanded,
+  setParentExpanded,
   initialExpanded = {},
   className,
   ariaLabel,
@@ -99,7 +112,12 @@ export const DataTable = ({
     sortable ? [{ id: sortCol ?? columns[0], desc: sortDir === 'desc' }] : [],
   );
   const [paging, setPaging] = React.useState<PaginationState>({ pageIndex, pageSize });
-  const [expanded, setExpanded] = React.useState<ExpandedState>(initialExpanded);
+  const [internalExpanded, setInternalExpanded] = React.useState<ExpandedState>(initialExpanded);
+
+  // Set which state value and setter function to use for tracking expanded rows.
+  // Use the parent state if passed in via props, otherwise track with internal state.
+  const expanded = parentExpanded ?? internalExpanded;
+  const setExpanded = setParentExpanded ?? setInternalExpanded;
 
   // Apply sorting to the full dataset first, then handle pagination
   const sortedData = React.useMemo(() => {