From 9732a3f6ac70ac75c6a3d1658debd4f76819c2b1 Mon Sep 17 00:00:00 2001
From: Maxim Dietz <maxim.dietz@goteleport.com>
Date: Wed, 13 Nov 2024 15:19:25 -0500
Subject: [PATCH] fix: Set initial data in `useTable` state

* Fix flickering between 'No data' state on first render, before
  `updateData` useEffect fires.
---
 web/packages/design/src/DataTable/useTable.ts | 95 ++++++++++++-------
 web/packages/design/src/utils/testing.tsx     | 11 +++
 .../src/JoinTokens/JoinTokens.test.tsx        |  9 +-
 3 files changed, 79 insertions(+), 36 deletions(-)

diff --git a/web/packages/design/src/DataTable/useTable.ts b/web/packages/design/src/DataTable/useTable.ts
index e8c70f3bd2dd0..16d69ee54f113 100644
--- a/web/packages/design/src/DataTable/useTable.ts
+++ b/web/packages/design/src/DataTable/useTable.ts
@@ -42,46 +42,71 @@ export default function useTable<T>({
   disableFilter = false,
   ...props
 }: TableProps<T>) {
-  const [state, setState] = useState<{
-    data: T[];
-    searchValue: string;
-    sort?: Sort<T>;
-    pagination?: Pagination<T>;
-  }>(() => {
+  // Determine the initial sort
+  let initialSort: Sort<T> | undefined;
+  if (!customSort) {
+    const { initialSort: initialSortProp } = props;
     // Finds the first sortable column to use for the initial sorting
     let col: TableColumn<T> | undefined;
-    if (!customSort) {
-      const { initialSort } = props;
-      if (initialSort) {
-        col = initialSort.altSortKey
-          ? columns.find(col => col.altSortKey === initialSort.altSortKey)
-          : columns.find(col => col.key === initialSort.key);
-      } else {
-        col = columns.find(column => column.isSortable);
-      }
+    if (initialSortProp) {
+      col = initialSortProp.altSortKey
+        ? columns.find(col => col.altSortKey === initialSortProp.altSortKey)
+        : columns.find(col => col.key === initialSortProp.key);
+    } else {
+      col = columns.find(column => column.isSortable);
     }
+    if (col) {
+      initialSort = {
+        key: (col.altSortKey || col.key) as keyof T,
+        onSort: col.onSort,
+        dir: initialSortProp?.dir || 'ASC',
+      };
+    }
+  }
 
-    return {
-      data: [],
-      searchValue: clientSearch?.initialSearchValue || '',
-      sort: col
-        ? {
-            key: (col.altSortKey || col.key) as keyof T,
-            onSort: col.onSort,
-            dir: props.initialSort?.dir || 'ASC',
-          }
-        : undefined,
-      pagination: pagination
-        ? {
-            paginatedData: paginateData([], pagination.pageSize),
-            currentPage: 0,
-            pagerPosition: pagination.pagerPosition,
-            pageSize: pagination.pageSize || 15,
-            CustomTable: pagination.CustomTable,
-          }
-        : undefined,
+  // Compute the initial data
+  const initialSearchValue = clientSearch?.initialSearchValue || '';
+  let initialData: T[];
+  if (serversideProps || disableFilter || !data?.length) {
+    initialData = data || [];
+  } else {
+    initialData = sortAndFilter(
+      data,
+      initialSearchValue,
+      initialSort,
+      searchableProps ||
+        (columns
+          .filter(column => column.key)
+          .map(column => column.key) as (keyof T & string)[]),
+      searchAndFilterCb,
+      showFirst
+    );
+  }
+
+  // Compute initial pagination if applicable
+  let initialPagination: Pagination<T> | undefined;
+  if (pagination) {
+    const pages = paginateData(initialData, pagination.pageSize);
+    initialPagination = {
+      paginatedData: pages,
+      currentPage: 0,
+      pagerPosition: pagination.pagerPosition,
+      pageSize: pagination.pageSize || 15,
+      CustomTable: pagination.CustomTable,
     };
-  });
+  }
+
+  const [state, setState] = useState<{
+    data: T[];
+    searchValue: string;
+    sort?: Sort<T>;
+    pagination?: Pagination<T>;
+  }>(() => ({
+    data: initialData,
+    searchValue: initialSearchValue,
+    sort: initialSort,
+    pagination: initialPagination,
+  }));
 
   function searchAndFilterCb(
     targetValue: any,
diff --git a/web/packages/design/src/utils/testing.tsx b/web/packages/design/src/utils/testing.tsx
index 2158692911414..6a133a97b605a 100644
--- a/web/packages/design/src/utils/testing.tsx
+++ b/web/packages/design/src/utils/testing.tsx
@@ -50,6 +50,16 @@ function render(
   return testingRender(ui, { wrapper: Providers, ...options });
 }
 
+/*
+ Returns a Promise resolving on the next macrotask, allowing any pending state
+ updates / timeouts to finish.
+ */
+function tick() {
+  return new Promise<void>(res =>
+    jest.requireActual('timers').setImmediate(res)
+  );
+}
+
 screen.debug = () => {
   window.console.log(prettyDOM());
 };
@@ -64,6 +74,7 @@ export {
   screen,
   fireEvent,
   darkTheme as theme,
+  tick,
   render,
   prettyDOM,
   waitFor,
diff --git a/web/packages/teleport/src/JoinTokens/JoinTokens.test.tsx b/web/packages/teleport/src/JoinTokens/JoinTokens.test.tsx
index 8b0962faabd7d..e4e1aacbb9391 100644
--- a/web/packages/teleport/src/JoinTokens/JoinTokens.test.tsx
+++ b/web/packages/teleport/src/JoinTokens/JoinTokens.test.tsx
@@ -15,7 +15,7 @@
  * You should have received a copy of the GNU Affero General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
-import { render, screen, fireEvent } from 'design/utils/testing';
+import { render, screen, fireEvent, act, tick } from 'design/utils/testing';
 import userEvent from '@testing-library/user-event';
 
 import selectEvent from 'react-select-event';
@@ -39,6 +39,13 @@ describe('JoinTokens', () => {
   test('edit dialog opens with values', async () => {
     const token = tokens[0];
     render(<Component />);
+
+    // DataTable re-renders before `userEvent.click` is fired, so `act(tick)`
+    // is used to wait for re-renders to complete.
+    // This wasn't an issue prior, as DataTable used to always mount with empty data,
+    // so `findAllByText` would wait a few ms before finding the text on commit 1.
+    await act(tick);
+
     const optionButtons = await screen.findAllByText(/options/i);
     await userEvent.click(optionButtons[0]);
     const editButtons = await screen.findAllByText(/view\/edit/i);