Avoid build timeout when having to wait for upgrade builds

node-src/lib/getEnv.ts

@@ -9,6 +9,7 @@ export interface Env {
   CHROMATIC_RETRIES: number;
   CHROMATIC_STORYBOOK_VERSION: string;
   CHROMATIC_TIMEOUT: number;
+  CHROMATIC_UPGRADE_TIMEOUT: number;
   ENVIRONMENT_WHITELIST: RegExp[];
   HTTP_PROXY: string;
   HTTPS_PROXY: string;
@@ -29,6 +30,7 @@ const {
   CHROMATIC_RETRIES = '5',
   CHROMATIC_STORYBOOK_VERSION,
   CHROMATIC_TIMEOUT = String(5 * 60 * 1000),
+  CHROMATIC_UPGRADE_TIMEOUT = String(60 * 60 * 1000),
   HTTP_PROXY = process.env.http_proxy,
   HTTPS_PROXY = process.env.https_proxy,
   LOGGLY_CUSTOMER_TOKEN = 'b5e26204-cdc5-4c78-a9cc-c69eb7fabad3',
@@ -65,6 +67,7 @@ export default () =>
     CHROMATIC_RETRIES: parseInt(CHROMATIC_RETRIES, 10),
     CHROMATIC_STORYBOOK_VERSION,
     CHROMATIC_TIMEOUT: parseInt(CHROMATIC_TIMEOUT, 10),
+    CHROMATIC_UPGRADE_TIMEOUT: parseInt(CHROMATIC_UPGRADE_TIMEOUT, 10),
     ENVIRONMENT_WHITELIST,
     HTTP_PROXY,
     HTTPS_PROXY,

node-src/lib/setExitCode.ts

@@ -36,6 +36,7 @@ export const exitCodes = {
   FETCH_ERROR: 201,
   GRAPHQL_ERROR: 202,
   MISSING_DEPENDENCY: 210,
+  VERIFICATION_TIMEOUT: 220,
   INVALID_OPTIONS: 254,
 };
 

node-src/tasks/verify.test.ts

@@ -1,8 +1,13 @@
 import { describe, expect, it, vi } from 'vitest';
 
+import { exitCodes } from '../lib/setExitCode';
 import { publishBuild, verifyBuild } from './verify';
 
-const env = { STORYBOOK_VERIFY_TIMEOUT: 1000 };
+const env = {
+  CHROMATIC_POLL_INTERVAL: 10,
+  CHROMATIC_UPGRADE_TIMEOUT: 100,
+  STORYBOOK_VERIFY_TIMEOUT: 20,
+};
 const log = { info: vi.fn(), warn: vi.fn(), debug: vi.fn() };
 const http = { fetch: vi.fn() };
 
@@ -55,9 +60,10 @@ describe('verifyBuild', () => {
       app: {},
       startedAt: Date.now(),
     };
-    const publishedBuild = { ...build, status: 'PUBLISHED', startedAt: null };
+    const publishedBuild = { ...build, status: 'PUBLISHED', startedAt: null, upgradeBuilds: [] };
     const client = { runQuery: vi.fn() };
     client.runQuery
+      // We can safely poll three times without hitting the timeout
       .mockReturnValueOnce({ app: { build: publishedBuild } })
       .mockReturnValueOnce({ app: { build: publishedBuild } })
       .mockReturnValue({ app: { build } });
@@ -95,6 +101,78 @@ describe('verifyBuild', () => {
     expect(ctx.skipSnapshots).toBe(undefined);
   });
 
+  it('times out if build takes too long to start', async () => {
+    const build = {
+      status: 'IN_PROGRESS',
+      features: { uiTests: true, uiReview: false },
+      app: {},
+      startedAt: Date.now(),
+    };
+    const publishedBuild = { ...build, status: 'PUBLISHED', startedAt: null, upgradeBuilds: [] };
+    const client = { runQuery: vi.fn() };
+    client.runQuery
+      // Polling four times is going to hit the timeout
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValue({ app: { build } });
+
+    const ctx = { client, ...defaultContext } as any;
+    await expect(verifyBuild(ctx, {} as any)).rejects.toThrow('Build verification timed out');
+
+    expect(client.runQuery).toHaveBeenCalledTimes(3);
+    expect(ctx.exitCode).toBe(exitCodes.VERIFICATION_TIMEOUT);
+  });
+
+  it('waits for upgrade builds before starting verification timeout', async () => {
+    const build = {
+      status: 'IN_PROGRESS',
+      features: { uiTests: true, uiReview: false },
+      app: {},
+      startedAt: Date.now(),
+    };
+    const upgradeBuilds = [{ completedAt: null }];
+    const completed = [{ completedAt: Date.now() }];
+    const publishedBuild = { ...build, status: 'PUBLISHED', startedAt: null, upgradeBuilds };
+    const client = { runQuery: vi.fn() };
+    client.runQuery
+      // Polling while upgrade builds are in progress is irrelevant
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      .mockReturnValueOnce({ app: { build: publishedBuild } })
+      // We can safely poll three times without hitting the timeout
+      .mockReturnValueOnce({ app: { build: { ...publishedBuild, upgradeBuilds: completed } } })
+      .mockReturnValueOnce({ app: { build: { ...publishedBuild, upgradeBuilds: completed } } })
+      .mockReturnValue({ app: { build } });
+
+    const ctx = { client, ...defaultContext } as any;
+    await verifyBuild(ctx, {} as any);
+
+    expect(ctx.build).toMatchObject(build);
+    expect(ctx.exitCode).toBe(undefined);
+  });
+
+  it('times out if upgrade builds take too long to complete', async () => {
+    const build = {
+      status: 'IN_PROGRESS',
+      features: { uiTests: true, uiReview: false },
+      app: {},
+      startedAt: Date.now(),
+    };
+    const upgradeBuilds = [{ completedAt: null }];
+    const publishedBuild = { ...build, status: 'PUBLISHED', startedAt: null, upgradeBuilds };
+    const client = { runQuery: vi.fn() };
+    client.runQuery.mockReturnValue({ app: { build: publishedBuild } });
+
+    const ctx = { client, ...defaultContext } as any;
+    await expect(verifyBuild(ctx, {} as any)).rejects.toThrow(
+      'Timed out waiting for upgrade builds to complete'
+    );
+  });
+
   it('sets exitCode to 5 if build was limited', async () => {
     const client = { runQuery: vi.fn() };
     client.runQuery.mockReturnValue({

node-src/tasks/verify.ts

@@ -12,6 +12,7 @@ import {
   pending,
   runOnlyFiles,
   runOnlyNames,
+  awaitingUpgrades,
   success,
   publishFailed,
 } from '../ui/tasks/verify';
@@ -85,15 +86,21 @@ const StartedBuildQuery = `
       build(number: $number) {
         startedAt
         failureReason
+        upgradeBuilds {
+          completedAt
+        }
       }
     }
   }
 `;
 interface StartedBuildQueryResult {
   app: {
     build: {
-      startedAt: number;
-      failureReason: string;
+      startedAt?: number;
+      failureReason?: string;
+      upgradeBuilds: {
+        completedAt?: number;
+      }[];
     };
   };
 }
@@ -178,6 +185,7 @@ export const verifyBuild = async (ctx: Context, task: Task) => {
     transitionTo(runOnlyNames)(ctx, task);
   }
 
+  let timeoutStart = Date.now();
   const waitForBuildToStart = async () => {
     const { storybookUrl } = ctx;
     const { number, reportToken } = ctx.announcedBuild;
@@ -187,12 +195,24 @@ export const verifyBuild = async (ctx: Context, task: Task) => {
     const {
       app: { build },
     } = await client.runQuery<StartedBuildQueryResult>(StartedBuildQuery, variables, options);
+
     if (build.failureReason) {
-      ctx.log.warn(brokenStorybook({ ...build, storybookUrl }));
+      ctx.log.warn(brokenStorybook({ failureReason: build.failureReason, storybookUrl }));
       setExitCode(ctx, exitCodes.STORYBOOK_BROKEN, true);
       throw new Error(publishFailed().output);
     }
+
     if (!build.startedAt) {
+      // Upgrade builds can take a long time to complete, so we can't apply a hard timeout yet,
+      // instead we only timeout on the actual build verification, after upgrades are complete.
+      if (build.upgradeBuilds.some((upgrade) => !upgrade.completedAt)) {
+        task.output = awaitingUpgrades(build).output;
+        timeoutStart = Date.now() + ctx.env.CHROMATIC_POLL_INTERVAL;
+      } else if (Date.now() - timeoutStart > ctx.env.STORYBOOK_VERIFY_TIMEOUT) {
+        setExitCode(ctx, exitCodes.VERIFICATION_TIMEOUT);
+        throw new Error('Build verification timed out');
+      }
+
       await delay(ctx.env.CHROMATIC_POLL_INTERVAL);
       await waitForBuildToStart();
       return;
@@ -209,8 +229,8 @@ export const verifyBuild = async (ctx: Context, task: Task) => {
     new Promise((_, reject) =>
       setTimeout(
         reject,
-        ctx.env.STORYBOOK_VERIFY_TIMEOUT,
-        new Error('Build verification timed out')
+        ctx.env.CHROMATIC_UPGRADE_TIMEOUT,
+        new Error('Timed out waiting for upgrade builds to complete')
       )
     ),
   ]);

node-src/ui/tasks/verify.ts

@@ -1,3 +1,5 @@
+import pluralize from 'pluralize';
+
 import { Context } from '../../types';
 
 export const initial = {
@@ -41,6 +43,20 @@ export const runOnlyNames = (ctx: Context) => ({
     .join(', ')}`,
 });
 
+export const awaitingUpgrades = ({
+  upgradeBuilds,
+}: {
+  upgradeBuilds: { completedAt?: number }[];
+}) => {
+  const pending = upgradeBuilds.filter((upgrade) => !upgrade.completedAt).length;
+  const upgrades = pluralize('upgrade build', upgradeBuilds.length, true);
+  return {
+    status: 'pending',
+    title: 'Verifying your Storybook',
+    output: `Waiting for ${pending}/${upgrades} to complete`,
+  };
+};
+
 export const success = (ctx: Context) => ({
   status: 'success',
   title: ctx.isPublishOnly ? `Published your Storybook` : `Started build ${ctx.build.number}`,