WIP: Addon-Vitest: Add test trigger API, take screenshots

code/addons/vitest/build-config.ts

@@ -12,6 +12,10 @@ const config: BuildEntries = {
         entryPoint: './src/manager.tsx',
         dts: false,
       },
+      {
+        exportEntries: ['./preview'],
+        entryPoint: './src/preview.ts',
+      },
       {
         exportEntries: ['./internal/setup-file'],
         entryPoint: './src/vitest-plugin/setup-file.ts',

code/addons/vitest/package.json

@@ -50,6 +50,10 @@
     "./package.json": "./package.json",
     "./postinstall": "./dist/postinstall.js",
     "./preset": "./dist/preset.js",
+    "./preview": {
+      "types": "./dist/preview.d.ts",
+      "default": "./dist/preview.js"
+    },
     "./vitest": "./dist/node/vitest.js",
     "./vitest-plugin": {
       "types": "./dist/vitest-plugin/index.d.ts",

code/addons/vitest/preview.js

@@ -0,0 +1 @@
+export * from './dist/preview.js';

code/addons/vitest/src/constants.ts

@@ -36,6 +36,8 @@ export const storeOptions = {
         coverage: false,
         a11y: false,
       },
+      componentTestStatuses: [],
+      a11yStatuses: [],
       componentTestCount: {
         success: 0,
         error: 0,
@@ -63,3 +65,4 @@ export const TEST_PROVIDER_STORE_CHANNEL_EVENT_NAME = 'UNIVERSAL_STORE:storybook
 
 export const STATUS_TYPE_ID_COMPONENT_TEST = 'storybook/component-test';
 export const STATUS_TYPE_ID_A11Y = 'storybook/a11y';
+export const STATUS_TYPE_ID_SCREENSHOT = 'storybook/screenshot';

code/addons/vitest/src/index.ts

@@ -1,3 +1,61 @@
-import { definePreviewAddon } from 'storybook/internal/csf';
+import { experimental_getStatusStore } from 'storybook/internal/core-server';
+import { type StoryId, definePreviewAddon } from 'storybook/internal/csf';
+
+import { STATUS_TYPE_ID_COMPONENT_TEST } from './constants';
+import type { Store } from './types';
 
 export default () => definePreviewAddon({});
+
+export async function triggerTestRun(actor: string, storyIds?: StoryId[]) {
+  // TODO: there must be a smarter way to share the store here
+  // while still lazy initializing it in the experimental_serverChannel preset
+  const store: Store | undefined = (globalThis as any).__STORYBOOK_ADDON_VITEST_STORE__;
+  if (!store) {
+    throw new Error('store not ready yet');
+  }
+
+  await store.untilReady();
+
+  const {
+    currentRun: { startedAt, finishedAt },
+  } = store.getState();
+  if (startedAt && !finishedAt) {
+    throw new Error('tests are already running');
+  }
+
+  store.send({
+    type: 'TRIGGER_RUN',
+    payload: {
+      storyIds,
+      triggeredBy: `external:${actor}`,
+    },
+  });
+
+  return new Promise((resolve, reject) => {
+    const unsubscribe = store.subscribe((event) => {
+      switch (event.type) {
+        case 'TEST_RUN_COMPLETED': {
+          console.log('Completed!');
+          console.dir(event.payload, { depth: 5 });
+          unsubscribe();
+          resolve(event.payload);
+          return;
+        }
+        case 'FATAL_ERROR': {
+          console.log('ERROR!');
+          console.dir(event.payload, { depth: 5 });
+          unsubscribe();
+          reject(event.payload);
+          return;
+        }
+        case 'CANCEL_RUN': {
+          console.log('CANCEL!');
+          console.dir(event, { depth: 5 });
+          unsubscribe();
+          reject('cancelled');
+          return;
+        }
+      }
+    });
+  });
+}

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

@@ -11,7 +11,12 @@ import type {
 import { throttle } from 'es-toolkit/function';
 import type { Report } from 'storybook/preview-api';
 
-import { STATUS_TYPE_ID_A11Y, STATUS_TYPE_ID_COMPONENT_TEST, storeOptions } from '../constants';
+import {
+  STATUS_TYPE_ID_A11Y,
+  STATUS_TYPE_ID_COMPONENT_TEST,
+  STATUS_TYPE_ID_SCREENSHOT,
+  storeOptions,
+} from '../constants';
 import type { RunTrigger, StoreEvent, StoreState, TriggerRunEvent, VitestError } from '../types';
 import { errorToErrorLike } from '../utils';
 import { VitestManager } from './vitest-manager';
@@ -21,6 +26,7 @@ export type TestManagerOptions = {
   store: experimental_UniversalStore<StoreState, StoreEvent>;
   componentTestStatusStore: StatusStoreByTypeId;
   a11yStatusStore: StatusStoreByTypeId;
+  screenshotStatusStore: StatusStoreByTypeId;
   testProviderStore: TestProviderStoreById;
   onError?: (message: string, error: Error) => void;
   onReady?: () => void;
@@ -43,6 +49,8 @@ export class TestManager {
 
   private a11yStatusStore: TestManagerOptions['a11yStatusStore'];
 
+  private screenshotStatusStore: TestManagerOptions['screenshotStatusStore'];
+
   private testProviderStore: TestManagerOptions['testProviderStore'];
 
   private onReady?: TestManagerOptions['onReady'];
@@ -59,6 +67,7 @@ export class TestManager {
     this.store = options.store;
     this.componentTestStatusStore = options.componentTestStatusStore;
     this.a11yStatusStore = options.a11yStatusStore;
+    this.screenshotStatusStore = options.screenshotStatusStore;
     this.testProviderStore = options.testProviderStore;
     this.onReady = options.onReady;
     this.storybookOptions = options.storybookOptions;
@@ -187,6 +196,36 @@ export class TestManager {
     const testCaseResultsToFlush = this.batchedTestCaseResults;
     this.batchedTestCaseResults = [];
 
+    const componentTestStatuses = testCaseResultsToFlush.map(({ storyId, testResult }) => ({
+      storyId,
+      typeId: STATUS_TYPE_ID_COMPONENT_TEST,
+      value: testStateToStatusValueMap[testResult.state],
+      title: 'Component tests',
+      description: testResult.errors?.map((error) => error.stack || error.message).join('\n') ?? '',
+      sidebarContextMenu: false,
+    }));
+
+    this.componentTestStatusStore.set(componentTestStatuses);
+
+    const a11yStatuses = testCaseResultsToFlush
+      .flatMap(({ storyId, reports }) =>
+        reports
+          ?.filter((r) => r.type === 'a11y')
+          .map((a11yReport) => ({
+            storyId,
+            typeId: STATUS_TYPE_ID_A11Y,
+            value: testStateToStatusValueMap[a11yReport.status],
+            title: 'Accessibility tests',
+            description: '',
+            sidebarContextMenu: false,
+          }))
+      )
+      .filter((a11yStatus) => a11yStatus !== undefined);
+
+    if (a11yStatuses.length > 0) {
+      this.a11yStatusStore.set(a11yStatuses);
+    }
+
     this.store.setState((s) => {
       let { success: ctSuccess, error: ctError } = s.currentRun.componentTestCount;
       let { success: a11ySuccess, warning: a11yWarning, error: a11yError } = s.currentRun.a11yCount;
@@ -216,6 +255,8 @@ export class TestManager {
           ...s.currentRun,
           componentTestCount: { success: ctSuccess, error: ctError },
           a11yCount: { success: a11ySuccess, warning: a11yWarning, error: a11yError },
+          componentTestStatuses: s.currentRun.componentTestStatuses.concat(componentTestStatuses),
+          a11yStatuses: s.currentRun.a11yStatuses.concat(a11yStatuses),
           // in some cases successes and errors can exceed the anticipated totalTestCount
           // e.g. when testing more tests than the stories we know about upfront
           // in those cases, we set the totalTestCount to the sum of successes and errors
@@ -226,36 +267,6 @@ export class TestManager {
         },
       };
     });
-
-    const componentTestStatuses = testCaseResultsToFlush.map(({ storyId, testResult }) => ({
-      storyId,
-      typeId: STATUS_TYPE_ID_COMPONENT_TEST,
-      value: testStateToStatusValueMap[testResult.state],
-      title: 'Component tests',
-      description: testResult.errors?.map((error) => error.stack || error.message).join('\n') ?? '',
-      sidebarContextMenu: false,
-    }));
-
-    this.componentTestStatusStore.set(componentTestStatuses);
-
-    const a11yStatuses = testCaseResultsToFlush
-      .flatMap(({ storyId, reports }) =>
-        reports
-          ?.filter((r) => r.type === 'a11y')
-          .map((a11yReport) => ({
-            storyId,
-            typeId: STATUS_TYPE_ID_A11Y,
-            value: testStateToStatusValueMap[a11yReport.status],
-            title: 'Accessibility tests',
-            description: '',
-            sidebarContextMenu: false,
-          }))
-      )
-      .filter((a11yStatus) => a11yStatus !== undefined);
-
-    if (a11yStatuses.length > 0) {
-      this.a11yStatusStore.set(a11yStatuses);
-    }
   }, 500);
 
   onTestRunEnd(endResult: { totalTestCount: number; unhandledErrors: VitestError[] }) {

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

@@ -11,6 +11,7 @@ import {
   ADDON_ID,
   STATUS_TYPE_ID_A11Y,
   STATUS_TYPE_ID_COMPONENT_TEST,
+  STATUS_TYPE_ID_SCREENSHOT,
   storeOptions,
 } from '../constants';
 import type { ErrorLike, FatalErrorEvent, StoreEvent, StoreState } from '../types';
@@ -41,6 +42,7 @@ new TestManager({
   store,
   componentTestStatusStore: getStatusStore(STATUS_TYPE_ID_COMPONENT_TEST),
   a11yStatusStore: getStatusStore(STATUS_TYPE_ID_A11Y),
+  screenshotStatusStore: getStatusStore(STATUS_TYPE_ID_SCREENSHOT),
   testProviderStore: getTestProviderStore(ADDON_ID),
   onReady: () => {
     process.send?.({ type: 'ready' });

code/addons/vitest/src/preset.ts

@@ -103,6 +103,7 @@ export const experimental_serverChannel = async (channel: Channel, options: Opti
       fsCache.set('state', selectCachedState(state));
     }
   });
+  globalThis.__STORYBOOK_ADDON_VITEST_STORE__ = store;
   const testProviderStore = experimental_getTestProviderStore(ADDON_ID);
 
   store.subscribe('TRIGGER_RUN', (event, eventInfo) => {

code/addons/vitest/src/preview.ts

@@ -0,0 +1,41 @@
+import type { StoryAnnotations } from 'storybook/internal/types';
+
+console.log('addon vitest preview!');
+
+// TODO: maybe this doesn't have to be an explicit preview annotation, but can be automatically added into vitest annotations somehow
+
+const preview: StoryAnnotations = {
+  afterEach: async ({ title, name, canvasElement, reporting }) => {
+    if (!(globalThis as any).__vitest_browser__) {
+      return;
+    }
+    //TODO: toggle this on an off based on something, probably like a11y
+    try {
+      console.log(`Taking screenshot for "${name}"`);
+      const { page } = await import('@vitest/browser/context');
+
+      const base64 = await page.screenshot({
+        path: `screenshots/${title}/${name}.png`,
+        base64: true,
+        element: canvasElement.firstChild,
+      });
+
+      reporting.addReport({
+        type: 'screenshot',
+        version: 1,
+        result: base64,
+        status: 'passed',
+      });
+    } catch (error) {
+      console.error('Error taking screenshot', error);
+      reporting.addReport({
+        type: 'screenshot',
+        version: 1,
+        result: error,
+        status: 'failed',
+      });
+    }
+  },
+};
+
+export default preview;

code/addons/vitest/src/types.ts

@@ -1,5 +1,5 @@
 import type { experimental_UniversalStore } from 'storybook/internal/core-server';
-import type { PreviewAnnotation, StoryId } from 'storybook/internal/types';
+import type { PreviewAnnotation, Status, StoryId } from 'storybook/internal/types';
 import type { API_HashEntry } from 'storybook/internal/types';
 
 export interface VitestError extends Error {
@@ -20,7 +20,12 @@ export type ErrorLike = {
   cause?: ErrorLike;
 };
 
-export type RunTrigger = 'run-all' | 'global' | 'watch' | Extract<API_HashEntry['type'], string>;
+export type RunTrigger =
+  | 'run-all'
+  | 'global'
+  | 'watch'
+  | Extract<API_HashEntry['type'], string>
+  | `external:${string}`;
 
 export type StoreState = {
   config: {
@@ -42,6 +47,8 @@ export type StoreState = {
   currentRun: {
     triggeredBy: RunTrigger | undefined;
     config: StoreState['config'];
+    componentTestStatuses: Status[];
+    a11yStatuses: Status[];
     componentTestCount: {
       success: number;
       error: number;