Lazy boot for Vitest runner

code/addons/test/package.json

@@ -80,6 +80,7 @@
     "@vitest/runner": "^2.1.1",
     "boxen": "^8.0.1",
     "find-up": "^7.0.0",
+    "lodash": "^4.17.21",
     "semver": "^7.6.3",
     "tinyrainbow": "^1.2.0",
     "ts-dedent": "^2.2.0",

code/addons/test/src/node/boot-test-runner.ts

@@ -11,64 +11,107 @@ import {
 
 import { log } from '../logger';
 
-export function bootTestRunner(channel: Channel) {
-  // This path is a bit confusing, but essentiall `boot-test-runner` gets bundled into the preset bundle
-  // which is at the root. Then, from the root, we want to load `node/vitest.js`
-  const sub = join(__dirname, 'node', 'vitest.js');
-
-  let child: ChildProcess;
-
-  function restartChildProcess() {
-    child?.kill();
-    log('Restarting Child Process...');
-    child = startChildProcess();
-  }
-
-  function startChildProcess() {
-    child = fork(sub, [], {
-      // We want to pipe output and error
-      // so that we can prefix the logs in the terminal
-      // with a clear identifier
-      stdio: ['inherit', 'pipe', 'pipe', 'ipc'],
-      silent: true,
-    });
+const MAX_RESTART_ATTEMPTS = 2;
 
-    child.stdout?.on('data', (data) => {
-      log(data);
-    });
+// This path is a bit confusing, but essentially `boot-test-runner` gets bundled into the preset bundle
+// which is at the root. Then, from the root, we want to load `node/vitest.js`
+const vitestModulePath = join(__dirname, 'node', 'vitest.js');
 
-    child.stderr?.on('data', (data) => {
-      log(data);
-    });
+export const bootTestRunner = (channel: Channel, initEvent?: string, initArgs?: any[]) =>
+  new Promise((resolve, reject) => {
+    let attempts = 0;
+    let child: null | ChildProcess;
 
-    child.on('message', (result: any) => {
-      if (result.type === 'error') {
-        log(result.message);
-        log(result.error);
-        restartChildProcess();
-      } else {
-        channel.emit(result.type, ...(result.args || []));
-      }
-    });
+    const forwardRun = (...args: any[]): void => {
+      child?.send({ type: TESTING_MODULE_RUN_REQUEST, args, from: 'server' });
+    };
+    const forwardRunAll = (...args: any[]): void => {
+      child?.send({ type: TESTING_MODULE_RUN_ALL_REQUEST, args, from: 'server' });
+    };
+    const forwardWatchMode = (...args: any[]): void => {
+      child?.send({ type: TESTING_MODULE_WATCH_MODE_REQUEST, args, from: 'server' });
+    };
+    const forwardCancel = (...args: any[]): void => {
+      child?.send({ type: TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, args, from: 'server' });
+    };
 
-    return child;
-  }
+    const startChildProcess = () => {
+      child = fork(vitestModulePath, [], {
+        // We want to pipe output and error
+        // so that we can prefix the logs in the terminal
+        // with a clear identifier
+        stdio: ['inherit', 'pipe', 'pipe', 'ipc'],
+        silent: true,
+      });
 
-  child = startChildProcess();
+      child.stdout?.on('data', (data) => {
+        log(data);
+      });
 
-  channel.on(TESTING_MODULE_RUN_REQUEST, (...args) => {
-    child.send({ type: TESTING_MODULE_RUN_REQUEST, args, from: 'server' });
-  });
+      child.stderr?.on('data', (data) => {
+        log(data);
+      });
 
-  channel.on(TESTING_MODULE_RUN_ALL_REQUEST, (...args) => {
-    child.send({ type: TESTING_MODULE_RUN_ALL_REQUEST, args, from: 'server' });
-  });
+      child.on('message', (result: any) => {
+        switch (result.type) {
+          case 'ready': {
+            attempts = 0;
+            child?.send({ type: initEvent, args: initArgs, from: 'server' });
+            channel.on(TESTING_MODULE_RUN_REQUEST, forwardRun);
+            channel.on(TESTING_MODULE_RUN_ALL_REQUEST, forwardRunAll);
+            channel.on(TESTING_MODULE_WATCH_MODE_REQUEST, forwardWatchMode);
+            channel.on(TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, forwardCancel);
+            channel.emit(result.type, ...(result.args || []));
+            resolve(result);
+            return;
+          }
 
-  channel.on(TESTING_MODULE_WATCH_MODE_REQUEST, (...args) => {
-    child.send({ type: TESTING_MODULE_WATCH_MODE_REQUEST, args, from: 'server' });
-  });
+          case 'error': {
+            channel.off(TESTING_MODULE_RUN_REQUEST, forwardRun);
+            channel.off(TESTING_MODULE_RUN_ALL_REQUEST, forwardRunAll);
+            channel.off(TESTING_MODULE_WATCH_MODE_REQUEST, forwardWatchMode);
+            channel.off(TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, forwardCancel);
 
-  channel.on(TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, (...args) => {
-    child.send({ type: TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, args, from: 'server' });
+            child?.kill();
+            child = null;
+
+            if (result.message) {
+              log(result.message);
+            }
+            if (result.error) {
+              log(result.error);
+            }
+
+            if (attempts >= MAX_RESTART_ATTEMPTS) {
+              log(`Aborting test runner process after ${MAX_RESTART_ATTEMPTS} restart attempts`);
+              channel.emit(
+                'error',
+                `Aborting test runner process after ${MAX_RESTART_ATTEMPTS} restart attempts`
+              );
+              reject(new Error('Test runner process failed to start'));
+            } else {
+              attempts += 1;
+              log(`Restarting test runner process (attempt ${attempts}/${MAX_RESTART_ATTEMPTS})`);
+              setTimeout(startChildProcess, 500);
+            }
+            return;
+          }
+        }
+      });
+    };
+
+    startChildProcess();
+
+    process.on('exit', () => {
+      child?.kill();
+      process.exit(0);
+    });
+    process.on('SIGINT', () => {
+      child?.kill();
+      process.exit(0);
+    });
+    process.on('SIGTERM', () => {
+      child?.kill();
+      process.exit(0);
+    });
   });
-}

code/addons/test/src/node/reporter.ts

@@ -4,6 +4,7 @@ import { type Reporter } from 'vitest/reporters';
 
 import type {
   TestingModuleRunAssertionResultPayload,
+  TestingModuleRunProgressPayload,
   TestingModuleRunResponsePayload,
   TestingModuleRunTestResultPayload,
 } from 'storybook/internal/core-events';
@@ -16,6 +17,7 @@ import type { Suite } from '@vitest/runner';
 // functions from the `@vitest/runner` package. It is not complex and does not have
 // any significant dependencies.
 import { getTests } from '@vitest/runner/utils';
+import throttle from 'lodash/throttle.js';
 
 import { TEST_PROVIDER_ID } from '../constants';
 import type { TestManager } from './test-manager';
@@ -42,7 +44,11 @@ export default class StorybookReporter implements Reporter {
 
   ctx!: Vitest;
 
-  constructor(private testManager: TestManager) {}
+  sendReport: (payload: TestingModuleRunProgressPayload) => void;
+
+  constructor(private testManager: TestManager) {
+    this.sendReport = throttle((payload) => this.testManager.sendProgressReport(payload), 200);
+  }
 
   onInit(ctx: Vitest) {
     this.ctx = ctx;
@@ -142,14 +148,14 @@ export default class StorybookReporter implements Reporter {
     try {
       const progress = this.getProgressReport();
 
-      this.testManager.sendProgressReport({
+      this.sendReport({
         status: 'success',
         payload: progress,
         providerId: TEST_PROVIDER_ID,
       });
     } catch (e) {
       if (e instanceof Error) {
-        this.testManager.sendProgressReport({
+        this.sendReport({
           status: 'failed',
           providerId: TEST_PROVIDER_ID,
           error: {
@@ -159,7 +165,7 @@ export default class StorybookReporter implements Reporter {
           },
         });
       } else {
-        this.testManager.sendProgressReport({
+        this.sendReport({
           status: 'failed',
           providerId: TEST_PROVIDER_ID,
           error: {

code/addons/test/src/node/test-manager.ts

@@ -33,6 +33,7 @@ export class TestManager {
   async restartVitest(watchMode = false) {
     await this.vitestManager.closeVitest();
     await this.vitestManager.startVitest(watchMode);
+    process.send?.({ type: 'ready', watchMode });
   }
 
   async handleWatchModeRequest(request: TestingModuleWatchModeRequestPayload) {

code/addons/test/src/node/vitest.ts

@@ -31,5 +31,18 @@ process.on('uncaughtException', (err) => {
 });
 
 process.on('unhandledRejection', (reason) => {
-  throw new Error(`Unhandled Rejection: ${reason}`);
+  throw reason;
+});
+
+process.on('exit', () => {
+  channel?.removeAllListeners();
+  process.exit(0);
+});
+process.on('SIGINT', () => {
+  channel?.removeAllListeners();
+  process.exit(0);
+});
+process.on('SIGTERM', () => {
+  channel?.removeAllListeners();
+  process.exit(0);
 });

code/addons/test/src/preset.ts

@@ -1,10 +1,47 @@
 import type { Channel } from 'storybook/internal/channels';
+import {
+  TESTING_MODULE_RUN_ALL_REQUEST,
+  TESTING_MODULE_RUN_REQUEST,
+  TESTING_MODULE_WATCH_MODE_REQUEST,
+} from 'storybook/internal/core-events';
 import type { Options } from 'storybook/internal/types';
 
 import { bootTestRunner } from './node/boot-test-runner';
 
 // eslint-disable-next-line @typescript-eslint/naming-convention
 export const experimental_serverChannel = async (channel: Channel, options: Options) => {
-  bootTestRunner(channel);
+  let booting = false;
+  let booted = false;
+  const start =
+    (eventName: string) =>
+    (...args: any[]) => {
+      if (!booted && !booting) {
+        booting = true;
+        bootTestRunner(channel, eventName, args)
+          .then(() => {
+            booted = true;
+          })
+          .catch(() => {
+            booted = false;
+          })
+          .finally(() => {
+            booting = false;
+          });
+      }
+    };
+
+  channel.on(TESTING_MODULE_RUN_ALL_REQUEST, start(TESTING_MODULE_RUN_ALL_REQUEST));
+  channel.on(TESTING_MODULE_RUN_REQUEST, start(TESTING_MODULE_RUN_REQUEST));
+  channel.on(TESTING_MODULE_WATCH_MODE_REQUEST, (payload) => {
+    if (payload.watchMode) {
+      start(TESTING_MODULE_WATCH_MODE_REQUEST)(payload);
+    }
+  });
+
   return channel;
 };
+
+// TODO:
+// 1 - Do not boot Vitest on Storybook boot, but rather on the first test run
+// 2 - Handle cases where Vitest is already booted, so we dont boot it again
+// 3 - Upon crash, provide a notification to the user

code/core/src/manager/components/sidebar/SidebarBottom.tsx

@@ -1,14 +1,20 @@
 import React, { useCallback, useEffect } from 'react';
 
 import { styled } from '@storybook/core/theming';
-import type { API_FilterFunction } from '@storybook/types';
+import type { API_FilterFunction, API_StatusUpdate, API_StatusValue } from '@storybook/types';
 
+import {
+  TESTING_MODULE_RUN_PROGRESS_RESPONSE,
+  type TestingModuleRunProgressPayload,
+  type TestingModuleRunResponsePayload,
+} from '@storybook/core/core-events';
 import {
   type API,
   type State,
   useStorybookApi,
   useStorybookState,
 } from '@storybook/core/manager-api';
+import { useChannel } from '@storybook/core/preview-api';
 
 import { FilterToggle } from './FilterToggle';
 
@@ -45,6 +51,30 @@ interface SidebarBottomProps {
   status: State['status'];
 }
 
+const statusMap: Record<any['status'], API_StatusValue> = {
+  failed: 'error',
+  passed: 'success',
+  pending: 'pending',
+};
+
+function processTestReport(payload: TestingModuleRunResponsePayload) {
+  const result: API_StatusUpdate = {};
+
+  payload.testResults.forEach((testResult: any) => {
+    testResult.results.forEach(({ storyId, status, failureMessages }: any) => {
+      if (storyId) {
+        result[storyId] = {
+          title: 'Vitest',
+          status: statusMap[status],
+          description: failureMessages?.length ? failureMessages.join('\n') : '',
+        };
+      }
+    });
+  });
+
+  return result;
+}
+
 export const SidebarBottomBase = ({ api, status = {} }: SidebarBottomProps) => {
   const [showWarnings, setShowWarnings] = React.useState(false);
   const [showErrors, setShowErrors] = React.useState(false);
@@ -99,5 +129,18 @@ export const SidebarBottomBase = ({ api, status = {} }: SidebarBottomProps) => {
 export const SidebarBottom = () => {
   const api = useStorybookApi();
   const { status } = useStorybookState();
+
+  useEffect(() => {
+    api.getChannel()?.on(TESTING_MODULE_RUN_PROGRESS_RESPONSE, (data) => {
+      if ('payload' in data) {
+        console.log('progress', data);
+        // TODO clear statuses
+        api.experimental_updateStatus('figure-out-id', processTestReport(data.payload));
+      } else {
+        console.log('error', data);
+      }
+    });
+  }, [api]);
+
   return <SidebarBottomBase api={api} status={status} />;
 };

code/yarn.lock

@@ -6246,6 +6246,7 @@ __metadata:
     boxen: "npm:^8.0.1"
     chalk: "npm:^5.3.0"
     find-up: "npm:^7.0.0"
+    lodash: "npm:^4.17.21"
     semver: "npm:^7.6.3"
     tinyrainbow: "npm:^1.2.0"
     ts-dedent: "npm:^2.2.0"