From 501bbddb6ec46b19f48877233dd3322fde765996 Mon Sep 17 00:00:00 2001
From: kripu77 <khadkakripu4@gmail.com>
Date: Mon, 6 Apr 2026 18:56:24 +1000
Subject: [PATCH] test: harden ci waits

---
 tests/components/diagram-preview.test.tsx | 16 ++++++++++------
 tests/e2e/grid-background.spec.ts         |  6 +++---
 tests/e2e/utils.ts                        | 18 +++++++++++++++---
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/tests/components/diagram-preview.test.tsx b/tests/components/diagram-preview.test.tsx
index e0afb37..a08428d 100644
--- a/tests/components/diagram-preview.test.tsx
+++ b/tests/components/diagram-preview.test.tsx
@@ -71,7 +71,9 @@ describe('DiagramPreview', () => {
       expect(screen.getByTestId('preview-board')).toBeInTheDocument();
     });
 
-    expect(updateViewport).toHaveBeenCalled();
+    await waitFor(() => {
+      expect(updateViewport).toHaveBeenCalled();
+    });
   });
 
   it('shows a fallback when mermaid parsing fails', async () => {
@@ -120,10 +122,12 @@ describe('DiagramPreview', () => {
       expect(value).toContain('1240');
     });
 
-    expect(updateViewport).toHaveBeenCalledWith(
-      mockBoard,
-      expect.any(Array),
-      expect.any(Number),
-    );
+    await waitFor(() => {
+      expect(updateViewport).toHaveBeenCalledWith(
+        mockBoard,
+        expect.any(Array),
+        expect.any(Number),
+      );
+    });
   });
 });
diff --git a/tests/e2e/grid-background.spec.ts b/tests/e2e/grid-background.spec.ts
index 17bf0a6..26b55f6 100644
--- a/tests/e2e/grid-background.spec.ts
+++ b/tests/e2e/grid-background.spec.ts
@@ -211,9 +211,9 @@ test.describe('Grid Background E2E Tests', () => {
       await page.reload();
       await waitForBoard(page);
       
-      const canvas = await getCanvas(page);
-      const content = await canvas.innerHTML();
-      expect(content).toContain('circle');
+      await expect
+        .poll(async () => (await getCanvas(page)).innerHTML(), { timeout: 10000 })
+        .toContain('circle');
     });
 
     test('should persist grid density setting', async ({ page }) => {
diff --git a/tests/e2e/utils.ts b/tests/e2e/utils.ts
index 831ba82..87d22fe 100644
--- a/tests/e2e/utils.ts
+++ b/tests/e2e/utils.ts
@@ -41,6 +41,12 @@ async function waitForBoardShell(page: Page): Promise<void> {
   await dismissOverlays(page);
 }
 
+async function waitForToolbar(page: Page): Promise<void> {
+  await expect(page.getByRole('button', { name: 'Select' }).first()).toBeVisible({
+    timeout: 30000,
+  });
+}
+
 export async function waitForMainBoard(page: Page): Promise<void> {
   let lastError: unknown;
 
@@ -186,6 +192,8 @@ export async function selectTool(page: Page, toolName: string): Promise<boolean>
     await page.waitForTimeout(100);
   }
 
+  await waitForToolbar(page);
+
   const directToolLabels: Record<string, string> = {
     arrow: 'Arrow',
     select: 'Select',
@@ -230,9 +238,10 @@ export async function selectTool(page: Page, toolName: string): Promise<boolean>
   if (toolName in shapeToolLabels) {
     const dropdownButton = page
       .getByTestId('shapes-dropdown-trigger')
+      .or(page.getByRole('button', { name: 'Shapes' }))
       .first();
 
-    if (!(await dropdownButton.isVisible({ timeout: 3000 }).catch(() => false))) {
+    if (!(await dropdownButton.isVisible({ timeout: 10000 }).catch(() => false))) {
       return false;
     }
 
@@ -264,15 +273,18 @@ export async function getElementCount(page: Page): Promise<number> {
   return page.evaluate(() => {
     const boardRoot = document.querySelector('[data-board="true"]');
     const count = boardRoot?.getAttribute('data-element-count');
+    const renderedElementCount = document.querySelectorAll(
+      '.board-wrapper [plait-data-id]',
+    ).length;
 
     if (count !== null && count !== undefined) {
       const parsed = Number.parseInt(count, 10);
       if (!Number.isNaN(parsed)) {
-        return parsed;
+        return Math.max(parsed, renderedElementCount);
       }
     }
 
-    return document.querySelectorAll('.board-wrapper [plait-data-id]').length;
+    return renderedElementCount;
   });
 }