test: harden background task completion waits (#382)

src/server/routes/api/accounts.add-background-task.test.ts

@@ -3,6 +3,7 @@ import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vites
 import { mkdtempSync, rmSync } from 'node:fs';
 import { tmpdir } from 'node:os';
 import { join } from 'node:path';
+import { waitForBackgroundTaskToReachTerminalState } from '../../test-fixtures/backgroundTaskTestUtils.js';
 
 const verifyTokenMock = vi.fn();
 const getApiTokensMock = vi.fn();
@@ -159,17 +160,16 @@ describe('accounts background initialization', () => {
 
       releaseTokens?.([{ name: 'default', value: 'sk-demo' }]);
 
-      for (let attempt = 0; attempt < 20; attempt += 1) {
-        const task = getBackgroundTask?.(body.jobId!);
-        if (task?.status === 'succeeded') break;
-        await new Promise((resolve) => setTimeout(resolve, 0));
-      }
+      const task = await waitForBackgroundTaskToReachTerminalState(
+        (taskId) => getBackgroundTask?.(taskId) ?? null,
+        body.jobId!,
+      );
 
       expect(syncTokensFromUpstreamMock).toHaveBeenCalledTimes(1);
       expect(refreshBalanceMock).toHaveBeenCalledTimes(1);
       expect(refreshModelsForAccountMock).toHaveBeenCalledTimes(1);
       expect(rebuildTokenRoutesFromAvailabilityMock).toHaveBeenCalledTimes(1);
-      expect(getBackgroundTask?.(body.jobId!)).toMatchObject({ status: 'succeeded' });
+      expect(task).toMatchObject({ status: 'succeeded' });
     } finally {
       releaseTokens?.([]);
       await responsePromise.catch(() => undefined);

src/server/routes/api/tokens.route-decision-refresh-task.test.ts

@@ -3,6 +3,7 @@ import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vites
 import { mkdtempSync, rmSync } from 'node:fs';
 import { tmpdir } from 'node:os';
 import { join } from 'node:path';
+import { waitForBackgroundTaskToReachTerminalState } from '../../test-fixtures/backgroundTaskTestUtils.js';
 
 type DbModule = typeof import('../../db/index.js');
 type BackgroundTaskModule = typeof import('../../services/backgroundTaskService.js');
@@ -163,13 +164,8 @@ describe('POST /api/routes/decision/refresh', () => {
     expect(body.jobId.length).toBeGreaterThan(10);
     expect(body.status).toBe('pending');
 
-    for (let attempt = 0; attempt < 200; attempt += 1) {
-      const task = getBackgroundTask(body.jobId);
-      if (task?.status === 'succeeded') break;
-      await new Promise((resolve) => setTimeout(resolve, 0));
-    }
-
-    expect(getBackgroundTask(body.jobId)).toMatchObject({ status: 'succeeded' });
+    const task = await waitForBackgroundTaskToReachTerminalState(getBackgroundTask, body.jobId);
+    expect(task).toMatchObject({ status: 'succeeded' });
 
     const routesResponse = await app.inject({
       method: 'GET',

src/server/routes/api/updateCenter.test.ts

@@ -4,6 +4,7 @@ import { mkdtempSync } from 'node:fs';
 import { join } from 'node:path';
 import { tmpdir } from 'node:os';
 import { eq } from 'drizzle-orm';
+import { waitForBackgroundTaskToReachTerminalState } from '../../test-fixtures/backgroundTaskTestUtils.js';
 
 const {
   fetchLatestStableGitHubReleaseMock,
@@ -801,11 +802,11 @@ describe('update center routes', () => {
 
     const deployBody = deployResponse.json() as { task: { id: string } };
 
-    for (let attempt = 0; attempt < 20; attempt += 1) {
-      const task = getBackgroundTask?.(deployBody.task.id);
-      if (task?.status === 'succeeded') break;
-      await new Promise((resolve) => setTimeout(resolve, 0));
-    }
+    const task = await waitForBackgroundTaskToReachTerminalState(
+      (taskId) => getBackgroundTask?.(taskId) ?? null,
+      deployBody.task.id,
+    );
+    expect(task).toMatchObject({ status: 'succeeded' });
 
     expect(getBackgroundTask?.(deployBody.task.id)?.logs).toEqual(expect.arrayContaining([
       expect.objectContaining({ message: 'Resolving target version' }),

src/server/test-fixtures/backgroundTaskTestUtils.ts

@@ -0,0 +1,31 @@
+import { type BackgroundTask } from '../services/backgroundTaskService.js';
+
+type WaitForBackgroundTaskOptions = {
+  timeoutMs?: number;
+  pollMs?: number;
+};
+
+const DEFAULT_TIMEOUT_MS = 5_000;
+const DEFAULT_POLL_MS = 10;
+
+function isBackgroundTaskTerminal(task: BackgroundTask | null): boolean {
+  return task?.status === 'succeeded' || task?.status === 'failed';
+}
+
+export async function waitForBackgroundTaskToReachTerminalState(
+  getTask: (taskId: string) => BackgroundTask | null,
+  taskId: string,
+  options: WaitForBackgroundTaskOptions = {},
+): Promise<BackgroundTask | null> {
+  const timeoutMs = Math.max(1, Math.trunc(options.timeoutMs ?? DEFAULT_TIMEOUT_MS));
+  const pollMs = Math.max(0, Math.trunc(options.pollMs ?? DEFAULT_POLL_MS));
+  const deadlineMs = Date.now() + timeoutMs;
+
+  let currentTask = getTask(taskId);
+  while (!isBackgroundTaskTerminal(currentTask) && Date.now() < deadlineMs) {
+    await new Promise((resolve) => setTimeout(resolve, pollMs));
+    currentTask = getTask(taskId);
+  }
+
+  return currentTask;
+}