From 64ae65950dd767e4b43d05031e183f480f2facf7 Mon Sep 17 00:00:00 2001
From: Stephen Haberman <shaberman@homebound.com>
Date: Thu, 12 Oct 2023 12:59:13 -0500
Subject: [PATCH] feat: Add column names to tableSnapshot.

---
 src/components/Table/GridTable.test.tsx | 27 +++++++++++++++++++++++++
 src/utils/rtl.tsx                       | 27 ++++++++++++++++---------
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/src/components/Table/GridTable.test.tsx b/src/components/Table/GridTable.test.tsx
index c615c5c2f..f4e1cc85e 100644
--- a/src/components/Table/GridTable.test.tsx
+++ b/src/components/Table/GridTable.test.tsx
@@ -32,6 +32,7 @@ import {
 type Data = { name: string; value: number | undefined | null };
 type Row = SimpleHeaderAndData<Data>;
 
+const idColumn: GridColumn<Row> = { id: "id", header: () => "Id", data: (data, { row }) => row.id };
 const nameColumn: GridColumn<Row> = { id: "name", header: () => "Name", data: ({ name }) => name };
 const valueColumn: GridColumn<Row> = { id: "value", header: () => "Value", data: ({ value }) => value };
 const columns = [nameColumn, valueColumn];
@@ -3265,6 +3266,32 @@ describe("GridTable", () => {
     `);
   });
 
+  it("tableSnapshot can use a subset of columns", async () => {
+    // Given a table with simple data
+    const r = await render(
+      <GridTable
+        columns={[idColumn, nameColumn, valueColumn]}
+        rows={[
+          simpleHeader,
+          { kind: "data", id: "1", data: { name: "Row 1", value: 200 } },
+          { kind: "data", id: "2", data: { name: "Row 2", value: 300 } },
+          { kind: "data", id: "3", data: { name: "Row 3", value: 1000 } },
+        ]}
+      />,
+    );
+
+    // Then a text snapshot should be generated when using `tableSnapshot`
+    expect(tableSnapshot(r, ["Id", "Value"])).toMatchInlineSnapshot(`
+      "
+      | Id | Value |
+      | -- | ----- |
+      | 1  | 200   |
+      | 2  | 300   |
+      | 3  | 1000  |
+      "
+    `);
+  });
+
   it("renders totals row in the correct order", async () => {
     type Row = SimpleHeaderAndData<Data> | TotalsRow;
     // Given a table with simple header, totals, and data row
diff --git a/src/utils/rtl.tsx b/src/utils/rtl.tsx
index 752a1fd8a..ad10738c8 100644
--- a/src/utils/rtl.tsx
+++ b/src/utils/rtl.tsx
@@ -109,25 +109,32 @@ export function rowAnd(r: RenderResult, rowNum: number, testId: string): HTMLEle
       "
     `);
  * */
-export function tableSnapshot(r: RenderResult): string {
+export function tableSnapshot(r: RenderResult, columnNames: string[] = []): string {
   const tableEl = r.getByTestId("gridTable");
   const dataRows = Array.from(tableEl.querySelectorAll("[data-gridrow]"));
   const hasExpandableHeader = !!tableEl.querySelector(`[data-testid="expandableColumn"]`);
 
-  const tableDataAsStrings = dataRows.map((row) => {
+  let tableDataAsStrings = dataRows.map((row) => {
     return Array.from(row.childNodes).map(getTextFromTableCellNode);
   });
 
-  return toMarkupTableString({ tableRows: tableDataAsStrings, hasExpandableHeader });
+  // If the user wants a subset of columns, look for column names
+  if (columnNames.length > 0) {
+    const headerCells = tableDataAsStrings[0];
+    if (headerCells) {
+      const columnIndices = columnNames.map((name) => {
+        const i = headerCells.indexOf(name);
+        if (i === -1) throw new Error(`Could not find header '${name}' in ${headerCells.join(", ")}`);
+        return i;
+      });
+      tableDataAsStrings = tableDataAsStrings.map((row) => columnIndices.map((index) => row[index]));
+    }
+  }
+
+  return toMarkupTableString(tableDataAsStrings, hasExpandableHeader);
 }
 
-function toMarkupTableString({
-  tableRows,
-  hasExpandableHeader,
-}: {
-  tableRows: (string | null)[][];
-  hasExpandableHeader: boolean;
-}) {
+function toMarkupTableString(tableRows: (string | null)[][], hasExpandableHeader: boolean): string {
   // Find the largest width of each column to set a consistent width for each row
   const columnWidths = tableRows.reduce((acc, row) => {
     row.forEach((cell, columnIndex) => {