Addon-Test: Implement Addon Test TestProvider Backend

code/.storybook/vitest.config.ts

@@ -31,9 +31,8 @@ export default mergeConfig(
     test: {
       name: 'storybook-ui',
       include: [
-        // TODO: test all core and addon stories later
-        // './core/**/components/**/*.{story,stories}.?(c|m)[jt]s?(x)',
-        '../addons/**/src/**/*.{story,stories}.?(c|m)[jt]s?(x)',
+        // TODO: Can be reverted. Temporarily I am adding all stories so that I can trigger tests for all stories in the UI.
+        '../{core,addons,lib}/**/{src,components,template}/**/*.{story,stories}.?(c|m)[jt]s?(x)',
       ],
       exclude: [
         ...defaultExclude,

code/addons/test/package.json

@@ -71,20 +71,23 @@
   },
   "dependencies": {
     "@storybook/csf": "^0.1.11",
-    "@storybook/icons": "^1.2.10"
+    "@storybook/icons": "^1.2.10",
+    "chalk": "^5.3.0"
   },
   "devDependencies": {
     "@types/semver": "^7",
-    "@vitest/browser": "^2.0.0",
+    "@vitest/browser": "^2.1.1",
+    "@vitest/runner": "^2.1.1",
     "boxen": "^8.0.1",
     "find-up": "^7.0.0",
     "semver": "^7.6.3",
     "tinyrainbow": "^1.2.0",
     "ts-dedent": "^2.2.0",
-    "vitest": "^2.0.0"
+    "vitest": "^2.1.1"
   },
   "peerDependencies": {
-    "@vitest/browser": "^2.0.0",
+    "@vitest/browser": "^2.1.1",
+    "@vitest/runner": "^2.1.1",
     "storybook": "workspace:^",
     "vitest": "^2.0.0"
   },
@@ -104,7 +107,8 @@
       "./src/preset.ts",
       "./src/vitest-plugin/index.ts",
       "./src/vitest-plugin/global-setup.ts",
-      "./src/postinstall.ts"
+      "./src/postinstall.ts",
+      "./src/node/vitest.ts"
     ]
   }
 }

code/addons/test/src/logger.ts

@@ -0,0 +1,7 @@
+import chalk from 'chalk';
+
+import { ADDON_ID } from './constants';
+
+export const log = (message: any) => {
+  console.log(`${chalk.magenta(ADDON_ID)}: ${message.toString().trim()}`);
+};

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

@@ -0,0 +1,74 @@
+import { type ChildProcess, fork } from 'node:child_process';
+import { join } from 'node:path';
+
+import type { Channel } from 'storybook/internal/channels';
+import {
+  TESTING_MODULE_CANCEL_TEST_RUN_REQUEST,
+  TESTING_MODULE_RUN_ALL_REQUEST,
+  TESTING_MODULE_RUN_REQUEST,
+  TESTING_MODULE_WATCH_MODE_REQUEST,
+} from 'storybook/internal/core-events';
+
+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,
+    });
+
+    child.stdout?.on('data', (data) => {
+      log(data);
+    });
+
+    child.stderr?.on('data', (data) => {
+      log(data);
+    });
+
+    child.on('message', (result: any) => {
+      if (result.type === 'error') {
+        log(result.message);
+        log(result.error);
+        restartChildProcess();
+      } else {
+        channel.emit(result.type, ...(result.args || []));
+      }
+    });
+
+    return child;
+  }
+
+  child = startChildProcess();
+
+  channel.on(TESTING_MODULE_RUN_REQUEST, (...args) => {
+    child.send({ type: TESTING_MODULE_RUN_REQUEST, args, from: 'server' });
+  });
+
+  channel.on(TESTING_MODULE_RUN_ALL_REQUEST, (...args) => {
+    child.send({ type: TESTING_MODULE_RUN_ALL_REQUEST, args, from: 'server' });
+  });
+
+  channel.on(TESTING_MODULE_WATCH_MODE_REQUEST, (...args) => {
+    child.send({ type: TESTING_MODULE_WATCH_MODE_REQUEST, args, from: 'server' });
+  });
+
+  channel.on(TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, (...args) => {
+    child.send({ type: TESTING_MODULE_CANCEL_TEST_RUN_REQUEST, args, from: 'server' });
+  });
+}

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

@@ -0,0 +1,191 @@
+import type { TaskState } from 'vitest';
+import type { Vitest } from 'vitest/node';
+import { type Reporter } from 'vitest/reporters';
+
+import type {
+  TestingModuleRunAssertionResultPayload,
+  TestingModuleRunResponsePayload,
+  TestingModuleRunTestResultPayload,
+} from 'storybook/internal/core-events';
+
+import type { API_StatusUpdate } from '@storybook/types';
+
+import type { Suite } from '@vitest/runner';
+// TODO
+// We can theoretically avoid the `@vitest/runner` dependency by copying over the necessary
+// functions from the `@vitest/runner` package. It is not complex and does not have
+// any significant dependencies.
+import { getTests } from '@vitest/runner/utils';
+
+import { TEST_PROVIDER_ID } from '../constants';
+import type { TestManager } from './test-manager';
+
+type Status = 'passed' | 'failed' | 'skipped' | 'pending' | 'todo' | 'disabled';
+
+function isDefined(value: any): value is NonNullable<typeof value> {
+  return value !== undefined && value !== null;
+}
+
+const StatusMap: Record<TaskState, Status> = {
+  fail: 'failed',
+  only: 'pending',
+  pass: 'passed',
+  run: 'pending',
+  skip: 'skipped',
+  todo: 'todo',
+};
+
+export default class StorybookReporter implements Reporter {
+  testStatusData: API_StatusUpdate = {};
+
+  start = 0;
+
+  ctx!: Vitest;
+
+  constructor(private testManager: TestManager) {}
+
+  onInit(ctx: Vitest) {
+    this.ctx = ctx;
+    this.start = Date.now();
+  }
+
+  getProgressReport(): TestingModuleRunResponsePayload {
+    const files = this.ctx.state.getFiles();
+    const fileTests = getTests(files);
+    // The number of total tests is dynamic and can change during the run
+    const numTotalTests = fileTests.length;
+
+    const numFailedTests = fileTests.filter((t) => t.result?.state === 'fail').length;
+    const numPassedTests = fileTests.filter((t) => t.result?.state === 'pass').length;
+    const numPendingTests = fileTests.filter(
+      (t) => t.result?.state === 'run' || t.mode === 'skip' || t.result?.state === 'skip'
+    ).length;
+    const testResults: Array<TestingModuleRunTestResultPayload> = [];
+
+    for (const file of files) {
+      const tests = getTests([file]);
+      let startTime = tests.reduce(
+        (prev, next) => Math.min(prev, next.result?.startTime ?? Number.POSITIVE_INFINITY),
+        Number.POSITIVE_INFINITY
+      );
+      if (startTime === Number.POSITIVE_INFINITY) {
+        startTime = this.start;
+      }
+
+      const endTime = tests.reduce(
+        (prev, next) =>
+          Math.max(prev, (next.result?.startTime ?? 0) + (next.result?.duration ?? 0)),
+        startTime
+      );
+
+      const assertionResults: TestingModuleRunAssertionResultPayload[] = tests
+        .map((t) => {
+          const ancestorTitles: string[] = [];
+          let iter: Suite | undefined = t.suite;
+          while (iter) {
+            ancestorTitles.push(iter.name);
+            iter = iter.suite;
+          }
+          ancestorTitles.reverse();
+
+          const status = StatusMap[t.result?.state || t.mode] || 'skipped';
+
+          if (status === 'passed' || status === 'pending') {
+            return {
+              status,
+              duration: t.result?.duration || 0,
+              storyId: (t.meta as any).storyId,
+            };
+          }
+
+          if (status === 'failed') {
+            return {
+              status,
+              duration: t.result?.duration || 0,
+              failureMessages: t.result?.errors?.map((e) => e.stack || e.message) || [],
+              storyId: (t.meta as any).storyId,
+            };
+          }
+
+          return null;
+        })
+        .filter(isDefined);
+
+      const hasFailedTests = tests.some((t) => t.result?.state === 'fail');
+
+      testResults.push({
+        results: assertionResults,
+        startTime,
+        endTime,
+        status: file.result?.state === 'fail' || hasFailedTests ? 'failed' : 'passed',
+        message: file.result?.errors?.[0]?.message,
+      });
+    }
+
+    return {
+      numFailedTests,
+      numPassedTests,
+      numPendingTests,
+      numTotalTests,
+      testResults,
+      success: true,
+      // TODO
+      // It is not simply (numPassedTests + numFailedTests) / numTotalTests
+      // because numTotalTests is dyanmic and can change during the run
+      // We need to calculate the progress based on the number of tests that have been run
+      progress: 0,
+      startTime: this.start,
+    };
+  }
+
+  async onTaskUpdate() {
+    const progress = this.getProgressReport();
+    try {
+      this.testManager.sendProgressReport({
+        status: 'success',
+        payload: progress,
+        providerId: TEST_PROVIDER_ID,
+      });
+    } catch (e) {
+      if (e instanceof Error) {
+        this.testManager.sendProgressReport({
+          status: 'failed',
+          providerId: TEST_PROVIDER_ID,
+          error: {
+            name: 'Failed to gather test results',
+            message: e.message,
+            stack: e.stack,
+          },
+        });
+      } else {
+        this.testManager.sendProgressReport({
+          status: 'failed',
+          providerId: TEST_PROVIDER_ID,
+          error: {
+            name: 'Failed to gather test results',
+            message: String(e),
+            stack: undefined,
+          },
+        });
+      }
+    }
+  }
+
+  // TODO
+  // Clearing the whole internal state of Vitest might be too aggressive
+  // Essentially, we want to reset the calculated total number of tests and the
+  // test results when a new test run starts, so that the getProgressReport
+  // method can calculate the correct values
+  async clearViteState() {
+    this.ctx.state.filesMap.clear();
+    this.ctx.state.pathsSet.clear();
+    this.ctx.state.idMap.clear();
+    this.ctx.state.errorsSet.clear();
+    this.ctx.state.processTimeoutCauses.clear();
+  }
+
+  async onFinished() {
+    this.clearViteState();
+  }
+}
+export { StorybookReporter };