fix(DatasourceControl): eliminate test flakiness and async race conditions (#35993)

Co-authored-by: Claude <noreply@anthropic.com>

superset-frontend/src/explore/components/controls/DatasourceControl/DatasourceControl.test.tsx

@@ -28,22 +28,43 @@ import {
   waitFor,
 } from 'spec/helpers/testing-library';
 import { fallbackExploreInitialData } from 'src/explore/fixtures';
+import type { DatasetObject, ColumnObject } from 'src/features/datasets/types';
 import DatasourceControl from '.';
 
 const SupersetClientGet = jest.spyOn(SupersetClient, 'get');
 
-afterEach(() => {
-  fetchMock.reset();
-  fetchMock.restore();
+let originalLocation: Location;
+
+beforeEach(() => {
+  originalLocation = window.location;
 });
 
-const mockDatasource = {
+afterEach(() => {
+  window.location = originalLocation;
+  fetchMock.reset();
+  fetchMock.restore();
+  jest.clearAllMocks(); // Clears mock history but keeps spy in place
+});
+
+type TestDatasource = Omit<
+  Partial<DatasetObject>,
+  'columns' | 'main_dttm_col'
+> & {
+  name: string;
+  database: { name: string };
+  columns?: Partial<ColumnObject>[];
+  type?: DatasourceType;
+  main_dttm_col?: string | null;
+};
+
+const mockDatasource: TestDatasource = {
   id: 25,
   database: {
     name: 'examples',
   },
   name: 'channels',
-  type: 'table',
+  datasource_name: 'channels',
+  type: DatasourceType.Table,
   columns: [],
   owners: [{ first_name: 'john', last_name: 'doe', id: 1, username: 'jd' }],
   sql: 'SELECT * FROM mock_datasource_sql',
@@ -80,7 +101,12 @@ const createProps = (overrides: JsonObject = {}) => ({
   ...overrides,
 });
 
-async function openAndSaveChanges(datasource: any) {
+async function openAndSaveChanges(datasource: TestDatasource) {
+  fetchMock.get(
+    'glob:*/api/v1/database/?q=*',
+    { result: [] },
+    { overwriteRoutes: true },
+  );
   fetchMock.put(
     'glob:*/api/v1/dataset/*',
     {},
@@ -95,10 +121,10 @@ async function openAndSaveChanges(datasource: any) {
       overwriteRoutes: true,
     },
   );
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
-  userEvent.click(await screen.findByTestId('edit-dataset'));
-  userEvent.click(await screen.findByTestId('datasource-modal-save'));
-  userEvent.click(await screen.findByText('OK'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(await screen.findByTestId('edit-dataset'));
+  await userEvent.click(await screen.findByTestId('datasource-modal-save'));
+  await userEvent.click(await screen.findByText('OK'));
 }
 
 test('Should render', async () => {
@@ -122,7 +148,7 @@ test('Should open a menu', async () => {
   expect(screen.queryByText('Swap dataset')).not.toBeInTheDocument();
   expect(screen.queryByText('View in SQL Lab')).not.toBeInTheDocument();
 
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   expect(await screen.findByText('Edit dataset')).toBeInTheDocument();
   expect(screen.getByText('Swap dataset')).toBeInTheDocument();
@@ -145,7 +171,7 @@ test('Should not show SQL Lab for non sql_lab role', async () => {
   });
   render(<DatasourceControl {...props} />, { useRouter: true });
 
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   expect(await screen.findByText('Edit dataset')).toBeInTheDocument();
   expect(screen.getByText('Swap dataset')).toBeInTheDocument();
@@ -168,7 +194,7 @@ test('Should show SQL Lab for sql_lab role', async () => {
   });
   render(<DatasourceControl {...props} />, { useRouter: true });
 
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   expect(await screen.findByText('Edit dataset')).toBeInTheDocument();
   expect(screen.getByText('Swap dataset')).toBeInTheDocument();
@@ -192,10 +218,10 @@ test('Click on Swap dataset option', async () => {
     useRedux: true,
     useRouter: true,
   });
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   await act(async () => {
-    userEvent.click(screen.getByText('Swap dataset'));
+    await userEvent.click(screen.getByText('Swap dataset'));
   });
   expect(
     screen.getByText(
@@ -209,14 +235,19 @@ test('Click on Edit dataset', async () => {
   SupersetClientGet.mockImplementationOnce(
     async () => ({ json: { result: [] } }) as any,
   );
+  fetchMock.get(
+    'glob:*/api/v1/database/?q=*',
+    { result: [] },
+    { overwriteRoutes: true },
+  );
   render(<DatasourceControl {...props} />, {
     useRedux: true,
     useRouter: true,
   });
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   await act(async () => {
-    userEvent.click(screen.getByText('Edit dataset'));
+    await userEvent.click(screen.getByText('Edit dataset'));
   });
 
   expect(
@@ -240,7 +271,7 @@ test('Edit dataset should be disabled when user is not admin', async () => {
     useRouter: true,
   });
 
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   expect(await screen.findByTestId('edit-dataset')).toHaveAttribute(
     'aria-disabled',
@@ -268,12 +299,12 @@ test('Click on View in SQL Lab', async () => {
       useRouter: true,
     },
   );
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   expect(queryByTestId('mock-sqllab-route')).not.toBeInTheDocument();
 
   await act(async () => {
-    userEvent.click(screen.getByText('View in SQL Lab'));
+    await userEvent.click(screen.getByText('View in SQL Lab'));
   });
 
   expect(getByTestId('mock-sqllab-route')).toBeInTheDocument();
@@ -302,7 +333,7 @@ test('Should open a different menu when datasource=query', async () => {
   expect(screen.queryByText('View in SQL Lab')).not.toBeInTheDocument();
   expect(screen.queryByText('Save as dataset')).not.toBeInTheDocument();
 
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
 
   expect(await screen.findByText('Query preview')).toBeInTheDocument();
   expect(screen.getByText('View in SQL Lab')).toBeInTheDocument();
@@ -323,7 +354,7 @@ test('Click on Save as dataset', async () => {
     useRedux: true,
     useRouter: true,
   });
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
   expect(
     screen.queryByRole('button', { name: /save/i }),
   ).not.toBeInTheDocument();
@@ -333,7 +364,7 @@ test('Click on Save as dataset', async () => {
   expect(
     screen.queryByText(/select or type dataset name/i),
   ).not.toBeInTheDocument();
-  userEvent.click(screen.getByText('Save as dataset'));
+  await userEvent.click(screen.getByText('Save as dataset'));
 
   // Renders a save dataset modal
   const saveRadioBtn = await screen.findByRole('radio', {
@@ -524,6 +555,12 @@ test('should allow creating new metrics in dataset editor', async () => {
   });
 
   // Mock API calls for dataset editor
+  fetchMock.get(
+    'glob:*/api/v1/database/?q=*',
+    { result: [] },
+    { overwriteRoutes: true },
+  );
+
   fetchMock.get(
     'glob:*/api/v1/dataset/*',
     { result: mockDatasourceWithMetrics },
@@ -551,40 +588,31 @@ test('should allow creating new metrics in dataset editor', async () => {
   });
 
   // Open datasource menu and click edit dataset
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
-  userEvent.click(await screen.findByTestId('edit-dataset'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(await screen.findByTestId('edit-dataset'));
 
   // Wait for modal to appear and navigate to Metrics tab
   await waitFor(() => {
     expect(screen.getByText('Metrics')).toBeInTheDocument();
   });
 
-  userEvent.click(screen.getByText('Metrics'));
+  await userEvent.click(screen.getByText('Metrics'));
 
   // Click add new metric button
-  await waitFor(() => {
-    const addButton = screen.getByTestId('crud-add-table-item');
-    expect(addButton).toBeInTheDocument();
-    userEvent.click(addButton);
-  });
+  const addButton = await screen.findByTestId('crud-add-table-item');
+  await userEvent.click(addButton);
 
   // Find and fill in the metric name
-  await waitFor(() => {
-    const nameInput = screen.getByTestId('textarea-editable-title-input');
-    expect(nameInput).toBeInTheDocument();
-    userEvent.clear(nameInput);
-    userEvent.type(nameInput, newMetricName);
-  });
+  const nameInput = await screen.findByTestId('textarea-editable-title-input');
+  await userEvent.clear(nameInput);
+  await userEvent.type(nameInput, newMetricName);
 
   // Save the modal
-  userEvent.click(screen.getByTestId('datasource-modal-save'));
+  await userEvent.click(screen.getByTestId('datasource-modal-save'));
 
   // Confirm the save
-  await waitFor(() => {
-    const okButton = screen.getByText('OK');
-    expect(okButton).toBeInTheDocument();
-    userEvent.click(okButton);
-  });
+  const okButton = await screen.findByText('OK');
+  await userEvent.click(okButton);
 
   // Verify the onDatasourceSave callback was called
   await waitFor(() => {
@@ -604,6 +632,12 @@ test('should allow deleting metrics in dataset editor', async () => {
   });
 
   // Mock API calls
+  fetchMock.get(
+    'glob:*/api/v1/database/?q=*',
+    { result: [] },
+    { overwriteRoutes: true },
+  );
+
   fetchMock.get(
     'glob:*/api/v1/dataset/*',
     { result: mockDatasourceWithMetrics },
@@ -626,36 +660,31 @@ test('should allow deleting metrics in dataset editor', async () => {
   });
 
   // Open edit dataset modal
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
-  userEvent.click(await screen.findByTestId('edit-dataset'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(await screen.findByTestId('edit-dataset'));
 
   // Navigate to Metrics tab
   await waitFor(() => {
     expect(screen.getByText('Metrics')).toBeInTheDocument();
   });
-  userEvent.click(screen.getByText('Metrics'));
+  await userEvent.click(screen.getByText('Metrics'));
 
   // Find existing metric and delete it
-  await waitFor(() => {
-    const metricRow = screen.getByText(existingMetricName).closest('tr');
-    expect(metricRow).toBeInTheDocument();
+  const metricRow = (await screen.findByText(existingMetricName)).closest('tr');
+  expect(metricRow).toBeInTheDocument();
 
-    const deleteButton = metricRow?.querySelector(
-      '[data-test="crud-delete-icon"]',
-    );
-    expect(deleteButton).toBeInTheDocument();
-    userEvent.click(deleteButton!);
-  });
+  const deleteButton = metricRow?.querySelector(
+    '[data-test="crud-delete-icon"]',
+  );
+  expect(deleteButton).toBeInTheDocument();
+  await userEvent.click(deleteButton!);
 
   // Save the changes
-  userEvent.click(screen.getByTestId('datasource-modal-save'));
+  await userEvent.click(screen.getByTestId('datasource-modal-save'));
 
   // Confirm the save
-  await waitFor(() => {
-    const okButton = screen.getByText('OK');
-    expect(okButton).toBeInTheDocument();
-    userEvent.click(okButton);
-  });
+  const okButton = await screen.findByText('OK');
+  await userEvent.click(okButton);
 
   // Verify the onDatasourceSave callback was called
   await waitFor(() => {
@@ -667,6 +696,12 @@ test('should handle metric save confirmation modal', async () => {
   const props = createProps();
 
   // Mock API calls for dataset editor
+  fetchMock.get(
+    'glob:*/api/v1/database/?q=*',
+    { result: [] },
+    { overwriteRoutes: true },
+  );
+
   fetchMock.get(
     'glob:*/api/v1/dataset/*',
     { result: mockDatasource },
@@ -689,15 +724,12 @@ test('should handle metric save confirmation modal', async () => {
   });
 
   // Open edit dataset modal
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
-  userEvent.click(await screen.findByTestId('edit-dataset'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(await screen.findByTestId('edit-dataset'));
 
   // Save without making changes
-  await waitFor(() => {
-    const saveButton = screen.getByTestId('datasource-modal-save');
-    expect(saveButton).toBeInTheDocument();
-    userEvent.click(saveButton);
-  });
+  const saveButton = await screen.findByTestId('datasource-modal-save');
+  await userEvent.click(saveButton);
 
   // Verify confirmation modal appears
   await waitFor(() => {
@@ -705,7 +737,7 @@ test('should handle metric save confirmation modal', async () => {
   });
 
   // Click OK to confirm
-  userEvent.click(screen.getByText('OK'));
+  await userEvent.click(screen.getByText('OK'));
 
   // Verify the save was processed
   await waitFor(() => {
@@ -724,6 +756,12 @@ test('should verify real DatasourceControl callback fires on save', async () =>
   });
 
   // Mock API calls with the same datasource (no changes needed for this test)
+  fetchMock.get(
+    'glob:*/api/v1/database/?q=*',
+    { result: [] },
+    { overwriteRoutes: true },
+  );
+
   fetchMock.get(
     'glob:*/api/v1/dataset/*',
     { result: mockDatasource },
@@ -750,8 +788,8 @@ test('should verify real DatasourceControl callback fires on save', async () =>
   expect(screen.getByTestId('datasource-control')).toBeInTheDocument();
 
   // Open dataset editor
-  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
-  userEvent.click(await screen.findByTestId('edit-dataset'));
+  await userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  await userEvent.click(await screen.findByTestId('edit-dataset'));
 
   // Wait for modal to open
   await waitFor(() => {
@@ -759,12 +797,9 @@ test('should verify real DatasourceControl callback fires on save', async () =>
   });
 
   // Save without making changes (this should still trigger the callback)
-  userEvent.click(screen.getByTestId('datasource-modal-save'));
-  await waitFor(() => {
-    const okButton = screen.getByText('OK');
-    expect(okButton).toBeInTheDocument();
-    userEvent.click(okButton);
-  });
+  await userEvent.click(screen.getByTestId('datasource-modal-save'));
+  const okButton = await screen.findByText('OK');
+  await userEvent.click(okButton);
 
   // Verify the REAL component called the callback
   // This tests that the integration point works (regardless of what data is passed)