test: fix CSV import Playwright test timeout by awaiting WebSocket completion before grid render (#27806)

Author: anuj-kumary

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/CSVImportWithQuotesAndCommas.spec.ts

@@ -19,6 +19,7 @@ import {
   validateImportStatus,
   waitForImportGridLoadMaskToDisappear,
 } from '../../utils/importUtils';
+import { setupCsvImportListener } from '../../utils/websocket';
 
 const cleanupTempFile = (filePath: string | undefined): void => {
   if (filePath && fs.existsSync(filePath)) {
@@ -40,7 +41,10 @@ const CSV_WITH_QUOTES_AND_COMMAS = `parent,name*,displayName,description,synonym
 ,"TermWithComma,AndQuote","Display name with ""quoted"" text, and comma","<p>Description with ""quotes"" and, commas</p>",,,,,,user:admin,Approved,,,`;
 
 test.describe('CSV Import with Commas and Quotes - All Entity Types', () => {
+  let createCsvImportPromise: () => Promise<void>;
+
   test.beforeEach(async ({ page }) => {
+    createCsvImportPromise = await setupCsvImportListener(page);
     await redirectToHomePage(page);
   });
 
@@ -67,6 +71,7 @@ test.describe('CSV Import with Commas and Quotes - All Entity Types', () => {
           {
             isContentString: true,
             tempFileName: `temp-quotes-commas-${uuid()}.csv`,
+            csvImportCompletedPromise: createCsvImportPromise(),
           }
         );
         tempFilePath = tempFile;
@@ -148,7 +153,9 @@ test.describe('CSV Import with Commas and Quotes - All Entity Types', () => {
       await page.click('[data-testid="import-button-description"]');
 
       try {
-        await uploadCSVAndWaitForGrid(page, exportedCsvPath);
+        await uploadCSVAndWaitForGrid(page, exportedCsvPath, {
+          csvImportCompletedPromise: createCsvImportPromise(),
+        });
 
         await expect(
           page.getByRole('gridcell', { name: 'Term1' }).first()

openmetadata-ui/src/main/resources/ui/playwright/utils/importUtils.ts

@@ -523,9 +523,11 @@ export const uploadCSVAndWaitForGrid = async (
   options?: {
     isContentString?: boolean;
     tempFileName?: string;
+    csvImportCompletedPromise?: Promise<void>;
   }
 ): Promise<{ rowCount: number; tempFilePath?: string }> => {
-  await page.locator('[type="file"]').waitFor({ state: 'attached' });
+  const uploadWidget = page.getByTestId('upload-file-widget');
+  await uploadWidget.waitFor({ state: 'attached' });
   let actualFilePath = filePath;
   let tempFilePath: string | undefined;
 
@@ -537,10 +539,11 @@ export const uploadCSVAndWaitForGrid = async (
     actualFilePath = tempFilePath;
   }
 
-  await page.setInputFiles('[type="file"]', actualFilePath);
-  await page
-    .getByTestId('upload-file-widget')
-    .waitFor({ state: 'hidden', timeout: 30000 });
+  await uploadWidget.setInputFiles([actualFilePath]);
+
+  if (options?.csvImportCompletedPromise) {
+    await options.csvImportCompletedPromise;
+  }
 
   await page
     .locator('.rdg-header-row')

openmetadata-ui/src/main/resources/ui/playwright/utils/websocket.ts

@@ -161,3 +161,30 @@ export const cleanupWebSocketMock = () => {
     defaultMock = null;
   }
 };
+
+export const setupCsvImportListener = async (
+  page: Page,
+  urlPattern = /push\/feed/
+): Promise<() => Promise<void>> => {
+  const pendingResolvers: Array<() => void> = [];
+
+  await page.routeWebSocket(urlPattern, (ws) => {
+    const server = ws.connectToServer();
+
+    server.onMessage((message) => {
+      ws.send(message);
+      const data =
+        typeof message === 'string' ? message : message.toString('utf-8');
+
+      if (data.includes('csvImportChannel') && data.includes('COMPLETED')) {
+        pendingResolvers.shift()?.();
+      }
+    });
+
+    ws.onMessage((message) => {
+      server.send(message);
+    });
+  });
+
+  return () => new Promise<void>((resolve) => pendingResolvers.push(resolve));
+};