From 1012f2f153605c36b5f72d4fabe4bc435f6390a1 Mon Sep 17 00:00:00 2001
From: spandev <143828642+Spandan-Mishra@users.noreply.github.com>
Date: Sun, 16 Nov 2025 03:31:01 +0530
Subject: [PATCH] fix: made show all columns work correctly (#25106)

* fix: made show all columns work correctly

* test: added test for show all columns in ColumnVisibilityButton

---------

Co-authored-by: Pallav <90088723+Pallava-Joshi@users.noreply.github.com>
---
 .../filters/ColumnVisibilityButton.test.tsx   | 55 +++++++++++++++++++
 .../filters/ColumnVisibilityButton.tsx        |  2 +-
 2 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 packages/features/data-table/components/filters/ColumnVisibilityButton.test.tsx

diff --git a/packages/features/data-table/components/filters/ColumnVisibilityButton.test.tsx b/packages/features/data-table/components/filters/ColumnVisibilityButton.test.tsx
new file mode 100644
index 0000000000..cc70193217
--- /dev/null
+++ b/packages/features/data-table/components/filters/ColumnVisibilityButton.test.tsx
@@ -0,0 +1,55 @@
+import { useEffect, useMemo } from "react";
+import { render, screen, fireEvent, waitFor } from "@testing-library/react";
+import { vi } from "vitest";
+import { useReactTable, getCoreRowModel, type Table, type Column } from "@tanstack/react-table";
+
+import { ColumnVisibilityButton } from "./ColumnVisibilityButton";
+
+let tableRef: Table<{ a: number; b: number; c: number }> | null = null;
+
+const TestWrapper = () => {
+  const data = useMemo(() => [{ a: 1, b: 2, c: 3 }], []);
+  const columns = useMemo(
+    () => [
+      { accessorKey: "a", header: "A" },
+      { accessorKey: "b", header: "B", enableHiding: true },
+      { accessorKey: "c", header: "C", enableHiding: true },
+    ],
+    []
+  );
+
+  const table = useReactTable({
+    data,
+    columns,
+    getCoreRowModel: getCoreRowModel(),
+    initialState: { columnVisibility: { b: false, c: false } },
+  });
+
+  useEffect(() => {
+    tableRef = table;
+  }, [table]);
+
+  return <ColumnVisibilityButton table={table} />;
+}
+
+describe("ColumnVisibilityButton", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    tableRef = null;
+  });
+
+  it("show_all_columns sets all columns visible", async () => {
+    render(<TestWrapper />);
+
+    fireEvent.click(screen.getByText("display"));
+
+    const showAll = await screen.findByText("show_all_columns");
+    fireEvent.click(showAll);
+
+    await waitFor(() => {
+      expect(tableRef).not.toBeNull();
+      const allVisible = tableRef!.getAllLeafColumns().every((col: Column<{ a: number; b: number; c: number }, unknown>) => col.getIsVisible());
+      expect(allVisible).toBe(true);
+    });
+  });
+});
diff --git a/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx b/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx
index c25894b97f..fc71d0e33e 100644
--- a/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx
+++ b/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx
@@ -78,7 +78,7 @@ function ColumnVisibilityButtonComponent<TData>(
           <CommandGroup>
             <CommandItem
               onSelect={() => {
-                allColumns.forEach((column) => column.toggleVisibility(true));
+                table.toggleAllColumnsVisible(true);
               }}
               className={classNames(
                 buttonClasses({ color: "secondary" }),