Test: Refactor to see stories as tests

code/addons/vitest/src/vitest-plugin/test-utils.ts

@@ -1,6 +1,6 @@
 import { type RunnerTask, type TaskMeta, type TestContext } from 'vitest';
 
-import { type Meta, type Story, isStory, toTestId } from 'storybook/internal/csf';
+import { type Meta, type Story, getStoryChildren, isStory, toTestId } from 'storybook/internal/csf';
 import type { ComponentAnnotations, ComposedStoryFn, Renderer } from 'storybook/internal/types';
 
 import { server } from '@vitest/browser/context';
@@ -41,8 +41,12 @@ export const testStory = (
   return async (context: TestContext & { story: ComposedStoryFn }) => {
     const annotations = getCsfFactoryAnnotations(story, meta);
 
-    const storyAnnotations =
-      isStory(story) && testName ? story.getAllTests()[testName].story.input : annotations.story;
+    const test =
+      isStory(story) && testName
+        ? getStoryChildren(story).find((child) => child.input.name === testName)
+        : undefined;
+
+    const storyAnnotations = test ? test.input : annotations.story;
 
     const composedStory = composeStory(
       storyAnnotations,
@@ -63,18 +67,15 @@ export const testStory = (
     };
 
     if (testName) {
+      // TODO: [test-syntax] isn't this done by the csf plugin somehow?
       _task.meta.storyId = toTestId(composedStory.id, testName);
     } else {
       _task.meta.storyId = composedStory.id;
     }
 
     await setViewport(composedStory.parameters, composedStory.globals);
 
-    if (isStory(story) && testName) {
-      await composedStory.run(undefined, story.getAllTests()[testName].test);
-    } else {
-      await composedStory.run(undefined);
-    }
+    await composedStory.run(undefined);
 
     _task.meta.reports = composedStory.reporting.reports;
   };

code/core/src/csf/csf-factories.test.ts

@@ -1,7 +1,7 @@
 //* @vitest-environment happy-dom */
 import { describe, expect, test, vi } from 'vitest';
 
-import { definePreview, definePreviewAddon } from './csf-factories';
+import { definePreview, definePreviewAddon, getStoryChildren } from './csf-factories';
 
 interface Addon1Types {
   parameters: { foo?: { value: string } };
@@ -53,11 +53,11 @@ describe('test function', () => {
 
     // register test
     MyStory.test(testName, testFn);
-    const { story: storyTestAnnotations } = MyStory.getAllTests()[testName];
-    expect(storyTestAnnotations.input.args).toEqual({ label: 'foo' });
+    const test = getStoryChildren(MyStory).find(({ input }) => input.name === testName)!;
+    expect(test.input.args).toEqual({ label: 'foo' });
 
     // execute test
-    await MyStory.run(undefined, testName);
+    await test.run(undefined, testName);
     expect(testFn).toHaveBeenCalled();
   });
   test('with overrides', async () => {
@@ -67,11 +67,11 @@ describe('test function', () => {
 
     // register test
     MyStory.test(testName, { args: { label: 'bar' } }, testFn);
-    const { story: storyTestAnnotations } = MyStory.getAllTests()[testName];
-    expect(storyTestAnnotations.input.args).toEqual({ label: 'bar' });
+    const test = getStoryChildren(MyStory).find(({ input }) => input.name === testName)!;
+    expect(test.input.args).toEqual({ label: 'bar' });
 
     // execute test
-    await MyStory.run(undefined, testName);
+    await test.run();
     expect(testFn).toHaveBeenCalled();
   });
 });

code/core/src/csf/csf-factories.ts

@@ -11,6 +11,8 @@ import type {
   TestFunction,
 } from 'storybook/internal/types';
 
+import { mountDestructured } from 'core/src/preview-api/modules/preview-web/render/mount-utils';
+
 import {
   combineParameters,
   composeConfigs,
@@ -154,7 +156,6 @@ export interface Story<
     annotations: StoryAnnotations<TRenderer, TRenderer['args']>,
     fn: TestFunction<TRenderer>
   ): void;
-  getAllTests(): Record<string, { story: Story<TRenderer>; test: TestFunction<TRenderer> }>;
 }
 
 export function isStory<TRenderer extends Renderer>(input: unknown): input is Story<TRenderer> {
@@ -178,14 +179,15 @@ function defineStory<
     return composed;
   };
 
-  const tests: Record<string, { story: Story<TRenderer>; test: TestFunction<TRenderer> }> = {};
+  const __children: Story<TRenderer>[] = [];
 
   return {
     _tag: 'Story',
     input,
     meta,
     // @ts-expect-error this is a private property used only once in renderers/react/src/preview
     __compose: compose,
+    __children,
     get composed() {
       const composed = compose();
       const { args, argTypes, parameters, id, tags, globals, storyName: name } = composed;
@@ -194,9 +196,8 @@ function defineStory<
     get play() {
       return input.play ?? meta.input?.play ?? (async () => {});
     },
-    async run(context, testName?: string) {
-      const composedRun = compose().run;
-      await composedRun(context, tests[testName!]?.test);
+    async run(context) {
+      await compose().run(context);
     },
     test(
       name: string,
@@ -206,10 +207,26 @@ function defineStory<
       const annotations = typeof overridesOrTestFn !== 'function' ? overridesOrTestFn : {};
       const testFunction = typeof overridesOrTestFn !== 'function' ? testFn! : overridesOrTestFn;
 
-      tests[name] = {
-        story: this.extend({ ...annotations, tags: [...(annotations.tags ?? []), 'test-fn'] }),
-        test: testFunction,
-      };
+      const play =
+        mountDestructured(this.play) || mountDestructured(testFn)
+          ? async ({ mount, context }: StoryContext<TRenderer>) => {
+              await this.play?.(context);
+              await testFunction(context);
+            }
+          : async (context: StoryContext<TRenderer>) => {
+              await this.play?.(context);
+              await testFunction(context);
+            };
+
+      const test = this.extend({
+        ...annotations,
+        name,
+        tags: ['test-fn', '!autodocs', ...(annotations.tags ?? [])],
+        play,
+      });
+      __children.push(test);
+
+      return test as unknown as void;
     },
     extend<TInput extends StoryAnnotations<TRenderer, TRenderer['args']>>(input: TInput) {
       return defineStory(
@@ -241,8 +258,14 @@ function defineStory<
         this.meta
       );
     },
-    getAllTests() {
-      return tests;
-    },
   };
 }
+
+export function getStoryChildren<TRenderer extends Renderer>(
+  story: Story<TRenderer>
+): Story<TRenderer>[] {
+  if ('__children' in story) {
+    return story.__children as Story<TRenderer>[];
+  }
+  return [];
+}