Merge pull request #6231 from Shopify/fix-flaky-poll-test

fix the flaky polling test in ci

Author: EvilGenius13

packages/app/src/cli/services/app-logs/logs-command/ui/components/hooks/usePollAppLogs.test.tsx

@@ -218,6 +218,9 @@ describe('usePollAppLogs', () => {
     // needed to await the render
     await vi.advanceTimersByTimeAsync(0)
 
+    // Wait for the async polling function to execute
+    await waitForMockCalls(mockedPollAppLogs, 1)
+
     expect(mockedPollAppLogs).toHaveBeenCalledTimes(1)
 
     expect(hook.lastResult?.appLogOutputs).toHaveLength(6)
@@ -340,6 +343,9 @@ describe('usePollAppLogs', () => {
     // needed to await the render
     await vi.advanceTimersByTimeAsync(0)
 
+    // Wait for the async polling function to execute
+    await waitForMockCalls(mockedPollAppLogs, 1)
+
     // Initial invocation, 401 returned
     expect(mockedPollAppLogs).toHaveBeenNthCalledWith(1, {
       pollOptions: {
@@ -354,6 +360,10 @@ describe('usePollAppLogs', () => {
 
     // Follow up invocation, which invokes resubscribeCallback
     await vi.advanceTimersToNextTimerAsync()
+
+    // Wait for the second async polling function call to execute
+    await waitForMockCalls(mockedPollAppLogs, 2)
+
     expect(mockedPollAppLogs).toHaveBeenNthCalledWith(2, {
       pollOptions: {jwtToken: NEW_JWT_TOKEN, cursor: '', filters: EMPTY_FILTERS},
       developerPlatformClient: mockedDeveloperPlatformClient,
@@ -525,3 +535,27 @@ function renderHook<THookResult>(renderHookCallback: () => THookResult) {
 
   return result
 }
+
+async function waitForMockCalls(mockFn: any, expectedCallCount: number, timeoutMs = 2000): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const timeout = setTimeout(
+      () =>
+        reject(
+          new Error(
+            `Mock timeout: expected ${expectedCallCount} calls, got ${mockFn.mock.calls.length} after ${timeoutMs}ms`,
+          ),
+        ),
+      timeoutMs,
+    )
+
+    const check = () => {
+      if (mockFn.mock.calls.length >= expectedCallCount) {
+        clearTimeout(timeout)
+        resolve()
+      } else {
+        setImmediate(check)
+      }
+    }
+    check()
+  })
+}