From bc62a98f51fd0405784cf536c81e9e710792860c Mon Sep 17 00:00:00 2001
From: Tymoteusz Czech <2625371+Tymek@users.noreply.github.com>
Date: Tue, 12 Dec 2023 14:01:04 +0100
Subject: [PATCH] update withTableState (#5603)

## About the changes
Handle column visibility from table state in URL and local storage.
---
 frontend/src/utils/withTableState.test.tsx | 307 +++++++++++++++++++++
 frontend/src/utils/withTableState.ts       |  64 ++++-
 2 files changed, 368 insertions(+), 3 deletions(-)
 create mode 100644 frontend/src/utils/withTableState.test.tsx

diff --git a/frontend/src/utils/withTableState.test.tsx b/frontend/src/utils/withTableState.test.tsx
new file mode 100644
index 0000000000..caeea0cf18
--- /dev/null
+++ b/frontend/src/utils/withTableState.test.tsx
@@ -0,0 +1,307 @@
+import { vi } from 'vitest';
+import { renderHook } from '@testing-library/react-hooks';
+import { useReactTable } from '@tanstack/react-table';
+import { withTableState } from './withTableState';
+import { useState } from 'react';
+import { render } from '@testing-library/react';
+
+describe('withTableState', () => {
+    it('should create paginated and sorted table state', () => {
+        const mockTableState = {
+            limit: 10,
+            offset: 10,
+            sortBy: 'name',
+            sortOrder: 'asc',
+        };
+        const mockSetTableState = vi.fn();
+        const mockOptions = { data: [], columns: [] };
+
+        const result = withTableState(
+            mockTableState,
+            mockSetTableState,
+            mockOptions,
+        );
+
+        expect(result.state).toEqual({
+            pagination: {
+                pageIndex: 1,
+                pageSize: 10,
+            },
+            sorting: [
+                {
+                    id: 'name',
+                    desc: false,
+                },
+            ],
+        });
+    });
+
+    it('sets default options', () => {
+        expect(
+            withTableState(
+                {
+                    limit: 10,
+                    offset: 10,
+                    sortBy: 'name',
+                    sortOrder: 'asc',
+                },
+                vi.fn(),
+                { data: [], columns: [] },
+            ),
+        ).toMatchObject({
+            getCoreRowModel: expect.any(Function),
+            enableSorting: true,
+            enableMultiSort: false,
+            manualPagination: true,
+            manualSorting: true,
+            enableSortingRemoval: false,
+            enableHiding: true,
+            onPaginationChange: expect.any(Function),
+            onSortingChange: expect.any(Function),
+            onColumnVisibilityChange: expect.any(Function),
+        });
+    });
+
+    it('should update page index and size', () => {
+        const mockTableState = {
+            limit: 10,
+            offset: 10,
+            sortBy: 'name',
+            sortOrder: 'asc',
+        };
+        const mockSetTableState = vi.fn();
+        const mockOptions = { data: [], columns: [] };
+
+        const { result } = renderHook(() =>
+            useReactTable(
+                withTableState(mockTableState, mockSetTableState, mockOptions),
+            ),
+        );
+
+        result.current.setPagination({
+            pageIndex: 3,
+            pageSize: 5,
+        });
+
+        expect(mockSetTableState).toHaveBeenCalledWith({
+            limit: 5,
+            offset: 15,
+        });
+    });
+
+    it('should update sorting', () => {
+        const mockTableState = {
+            limit: 10,
+            offset: 10,
+            sortBy: 'name',
+            sortOrder: 'asc',
+        };
+        const mockSetTableState = vi.fn();
+        const mockOptions = { data: [], columns: [] };
+
+        const { result } = renderHook(() =>
+            useReactTable(
+                withTableState(mockTableState, mockSetTableState, mockOptions),
+            ),
+        );
+
+        result.current.setSorting([
+            {
+                id: 'createdAt',
+                desc: true,
+            },
+        ]);
+
+        expect(mockSetTableState).toHaveBeenCalledWith({
+            sortBy: 'createdAt',
+            sortOrder: 'desc',
+        });
+    });
+
+    it('should handle column visibility', () => {
+        const mockTableState = {
+            limit: 10,
+            offset: 10,
+            sortBy: 'name',
+            sortOrder: 'asc',
+            columns: ['name'],
+        };
+        const mockSetTableState = vi.fn();
+        const mockOptions = {
+            data: [],
+            columns: [
+                {
+                    id: 'name',
+                    show: true,
+                },
+                {
+                    id: 'createdAt',
+                    show: false,
+                },
+            ],
+        };
+
+        const { result } = renderHook(() =>
+            useReactTable(
+                withTableState(mockTableState, mockSetTableState, mockOptions),
+            ),
+        );
+
+        expect(result.current.getState().columnVisibility).toMatchObject({
+            name: true,
+        });
+
+        result.current.setColumnVisibility({ name: false, createdAt: true });
+
+        expect(mockSetTableState).toHaveBeenCalledWith({
+            columns: ['createdAt'],
+        });
+    });
+
+    it('is always using external state', () => {
+        const initialProps = {
+            limit: 5,
+            offset: 40,
+            sortBy: 'name',
+            sortOrder: 'desc',
+        };
+
+        const { result, rerender } = renderHook(
+            (state) =>
+                useReactTable(
+                    withTableState(state as any, vi.fn(), {
+                        data: [],
+                        columns: [],
+                    }),
+                ),
+            { initialProps },
+        );
+
+        expect(result.current.getState()).toMatchObject({
+            pagination: {
+                pageIndex: 8,
+                pageSize: 5,
+            },
+            sorting: [
+                {
+                    id: 'name',
+                    desc: true,
+                },
+            ],
+        });
+
+        rerender({
+            limit: 10,
+            offset: 10,
+            sortBy: 'createdAt',
+            sortOrder: 'asc',
+        });
+
+        expect(result.current.getState()).toMatchObject({
+            pagination: {
+                pageIndex: 1,
+                pageSize: 10,
+            },
+            sorting: [
+                {
+                    id: 'createdAt',
+                    desc: false,
+                },
+            ],
+        });
+    });
+
+    it('works end-to-end with useReactTable', () => {
+        const Component = () => {
+            const [state, setState] = useState({
+                limit: 5,
+                offset: 40,
+                sortBy: 'name',
+                sortOrder: 'desc',
+            });
+
+            const setTableState = (newState: any) => {
+                setState((state) => ({ ...state, ...newState }));
+            };
+
+            const table = useReactTable(
+                withTableState(state, setTableState, {
+                    data: [],
+                    columns: [],
+                }),
+            );
+
+            return (
+                <>
+                    <button type='button' onClick={table.nextPage}>
+                        Next page
+                    </button>
+                    <button type='button' onClick={table.previousPage}>
+                        Previous page
+                    </button>
+                    <button
+                        type='button'
+                        onClick={() =>
+                            table.setPagination({ pageIndex: 2, pageSize: 10 })
+                        }
+                    >
+                        Paginate
+                    </button>
+                    <button
+                        type='button'
+                        onClick={() =>
+                            table.setSorting([{ id: 'createdAt', desc: true }])
+                        }
+                    >
+                        Sort
+                    </button>
+                    <textarea
+                        value={JSON.stringify(table.getState())}
+                        readOnly
+                    />
+                    <input
+                        data-testid='page'
+                        type='text'
+                        value={table.getState().pagination.pageIndex}
+                        readOnly
+                    />
+                    <input
+                        data-testid='pageSize'
+                        type='text'
+                        value={table.getState().pagination.pageSize}
+                        readOnly
+                    />
+                    <input
+                        data-testid='sort'
+                        type='text'
+                        value={table.getState().sorting[0].id}
+                        readOnly
+                    />
+                </>
+            );
+        };
+
+        const { getByTestId, getByRole } = render(<Component />);
+
+        expect(getByTestId('page')).toHaveValue('8');
+        expect(getByTestId('pageSize')).toHaveValue('5');
+        expect(getByTestId('sort')).toHaveValue('name');
+
+        getByRole('button', { name: 'Next page' }).click();
+
+        expect(getByTestId('page')).toHaveValue('9');
+
+        getByRole('button', { name: 'Previous page' }).click();
+
+        expect(getByTestId('page')).toHaveValue('8');
+
+        getByRole('button', { name: 'Paginate' }).click();
+
+        expect(getByTestId('page')).toHaveValue('2');
+        expect(getByTestId('pageSize')).toHaveValue('10');
+
+        getByRole('button', { name: 'Sort' }).click();
+
+        expect(getByTestId('sort')).toHaveValue('createdAt');
+    });
+});
diff --git a/frontend/src/utils/withTableState.ts b/frontend/src/utils/withTableState.ts
index 7c9a574996..6a221e8996 100644
--- a/frontend/src/utils/withTableState.ts
+++ b/frontend/src/utils/withTableState.ts
@@ -3,9 +3,12 @@ import {
     type SortingState,
     type PaginationState,
     type TableOptions,
+    type VisibilityState,
     getCoreRowModel,
 } from '@tanstack/react-table';
 
+type TableStateColumnsType = (string | null)[] | null;
+
 const createOnSortingChange =
     (
         tableState: {
@@ -73,6 +76,39 @@ const createOnPaginationChange =
         }
     };
 
+const createOnColumnVisibilityChange =
+    (
+        tableState: {
+            columns?: TableStateColumnsType;
+        },
+        setTableState: (newState: {
+            columns?: TableStateColumnsType;
+        }) => void,
+    ): OnChangeFn<VisibilityState> =>
+    (newVisibility) => {
+        const columnsObject = tableState.columns?.reduce(
+            (acc, column) => ({
+                ...acc,
+                ...(column && { [column]: true }),
+            }),
+            {},
+        );
+
+        if (typeof newVisibility === 'function') {
+            const computedVisibility = newVisibility(columnsObject || {});
+            const columns = Object.keys(computedVisibility).filter(
+                (column) => computedVisibility[column],
+            );
+
+            setTableState({ columns });
+        } else {
+            const columns = Object.keys(newVisibility).filter(
+                (column) => newVisibility[column],
+            );
+            setTableState({ columns });
+        }
+    };
+
 const createSortingState = (tableState: {
     sortBy: string;
     sortOrder: string;
@@ -95,18 +131,35 @@ const createPaginationState = (tableState: {
     },
 });
 
+const createColumnVisibilityState = (tableState: {
+    columns?: TableStateColumnsType;
+}) =>
+    tableState.columns
+        ? {
+              columnVisibility: tableState.columns?.reduce(
+                  (acc, column) => ({
+                      ...acc,
+                      ...(column && { [column]: true }),
+                  }),
+                  {},
+              ),
+          }
+        : {};
+
 export const withTableState = <T extends Object>(
     tableState: {
         sortBy: string;
         sortOrder: string;
         limit: number;
         offset: number;
+        columns?: TableStateColumnsType;
     },
     setTableState: (newState: {
         sortBy?: string;
         sortOrder?: string;
         limit?: number;
         offset?: number;
+        columns?: TableStateColumnsType;
     }) => void,
     options: Omit<TableOptions<T>, 'getCoreRowModel'>,
 ) => ({
@@ -117,12 +170,17 @@ export const withTableState = <T extends Object>(
     manualSorting: true,
     enableSortingRemoval: false,
     enableHiding: true,
+    onPaginationChange: createOnPaginationChange(tableState, setTableState),
+    onSortingChange: createOnSortingChange(tableState, setTableState),
+    onColumnVisibilityChange: createOnColumnVisibilityChange(
+        tableState,
+        setTableState,
+    ),
+    ...options,
     state: {
         ...createSortingState(tableState),
         ...createPaginationState(tableState),
+        ...createColumnVisibilityState(tableState),
         ...(options.state || {}),
     },
-    onPaginationChange: createOnPaginationChange(tableState, setTableState),
-    onSortingChange: createOnSortingChange(tableState, setTableState),
-    ...options,
 });