fix(DataGrid): track action columns correctly (#828)

## 📝 Changes

- fix DataGrid using row count instead of column count to track where
columns were in the table
- Adds story and test

Bug is causing gaps in tables with 1 row:

<img width="978" alt="image"
src="https://github.com/EasyPost/easy-ui/assets/752942/e854726e-2d21-4826-b936-1cdd01c8ed8a">

## ✅ Checklist

- [x] Unit tests are written and passing
- [x] Changeset is added

.changeset/six-turkeys-speak.md

@@ -0,0 +1,5 @@
+---
+"@easypost/easy-ui": patch
+---
+
+fix(DataGrid): track action columns correctly

easy-ui-react/src/DataGrid/Cell.tsx

@@ -33,7 +33,7 @@ export function Cell({ cell, state }: CellProps) {
   const hasLeftShadow =
     hasActionsAtEnd &&
     table.isRightEdgeUnderScroll &&
-    cell.index === state.collection.size - 1;
+    cell.index === state.collection.columnCount - 1;
 
   const className = classNames(
     styles.Cell,
@@ -44,12 +44,12 @@ export function Cell({ cell, state }: CellProps) {
     cell.index === 0 && styles.first,
     hasActionsAtStart && cell.index === 0 && styles.firstWithActions,
     hasActionsAtStart && cell.index === 1 && styles.secondWithActions,
-    cell.index === state.collection.size - 1 && styles.last,
+    cell.index === state.collection.columnCount - 1 && styles.last,
     hasActionsAtEnd &&
-      cell.index === state.collection.size - 1 &&
+      cell.index === state.collection.columnCount - 1 &&
       styles.lastWithActions,
     hasActionsAtEnd &&
-      cell.index === state.collection.size - 2 &&
+      cell.index === state.collection.columnCount - 2 &&
       styles.secondToLastWithActions,
   );
 

easy-ui-react/src/DataGrid/ColumnHeader.tsx

@@ -38,7 +38,7 @@ export function ColumnHeader({ column, state }: ColumnHeaderProps) {
   const hasLeftShadow =
     hasActionsAtEnd &&
     table.isRightEdgeUnderScroll &&
-    column.index === state.collection.size - 1;
+    column.index === state.collection.columnCount - 1;
 
   const className = classNames(
     styles.ColumnHeader,
@@ -50,12 +50,12 @@ export function ColumnHeader({ column, state }: ColumnHeaderProps) {
     column.index === 0 && styles.first,
     hasActionsAtStart && column.index === 0 && styles.firstWithActions,
     hasActionsAtStart && column.index === 1 && styles.secondWithActions,
-    column.index === state.collection.size - 1 && styles.last,
+    column.index === state.collection.columnCount - 1 && styles.last,
     hasActionsAtEnd &&
-      column.index === state.collection.size - 1 &&
+      column.index === state.collection.columnCount - 1 &&
       styles.lastWithActions,
     hasActionsAtEnd &&
-      column.index === state.collection.size - 2 &&
+      column.index === state.collection.columnCount - 2 &&
       styles.secondToLastWithActions,
   );
 

easy-ui-react/src/DataGrid/DataGrid.stories.tsx

@@ -227,6 +227,30 @@ export const WithIcons: Story = {
   },
 };
 
+export const WithRowExpansionAndKebabMenu: Story = {
+  render: Template.bind({}),
+  args: {
+    "aria-label": "Example data grid with row expansion and kebab menu",
+    rows,
+    renderExpandedRow: (rowKey: Key) => (
+      <PlaceholderBox width="100%" height="140px">
+        Space for row {rowKey} content
+      </PlaceholderBox>
+    ),
+    rowActions: () => [
+      {
+        type: "menu",
+        renderMenuOverlay: () => (
+          <Menu.Overlay onAction={action("Menu item clicked!")}>
+            <Menu.Item>Action 1</Menu.Item>
+            <Menu.Item>Action 2</Menu.Item>
+          </Menu.Overlay>
+        ),
+      },
+    ],
+  },
+};
+
 function WithSortTemplate(args: Partial<DataGridProps>) {
   // https://react-spectrum.adobe.com/react-stately/useAsyncList.html
   const list = useAsyncList({

easy-ui-react/src/DataGrid/DataGrid.test.tsx

@@ -178,6 +178,38 @@ describe("<DataGrid />", () => {
       direction: "ascending",
     });
   });
+
+  it("should find row expansion and kebab menu cells", async () => {
+    render(
+      createDataGrid({
+        renderExpandedRow: (rowKey: Key) => <div>Row {rowKey} content</div>,
+        rowActions: () => [
+          {
+            type: "menu",
+            renderMenuOverlay: () => (
+              <Menu.Overlay>
+                <Menu.Item>Action 1</Menu.Item>
+              </Menu.Overlay>
+            ),
+          },
+        ],
+      }),
+    );
+
+    const row = getRow(1);
+    const cells = [...row.children];
+    const first = cells[0];
+    const last = cells[cells.length - 1];
+
+    expect(first).toHaveAttribute(
+      "class",
+      expect.stringContaining("firstWithActions"),
+    );
+    expect(last).toHaveAttribute(
+      "class",
+      expect.stringContaining("lastWithActions"),
+    );
+  });
 });
 
 const columns = [