Test: Refactor .test prototype

code/.storybook/preview.tsx

@@ -215,18 +215,21 @@ const decorators = [
    * This decorator renders the stories side-by-side, stacked or default based on the theme switcher
    * in the toolbar
    */
-  (StoryFn, { globals, playFunction, args, storyGlobals, parameters }) => {
+  (StoryFn, { globals, playFunction, testFunction, args, storyGlobals, parameters }) => {
     let theme = globals.sb_theme;
     let showPlayFnNotice = false;
 
     // this makes the decorator be out of 'phase' with the actually selected theme in the toolbar
     // but this is acceptable, I guess
     // we need to ensure only a single rendering in chromatic
     // a more 'correct' approach would be to set a specific theme global on every story that has a playFunction
-    if (playFunction && args.autoplay !== false && !(theme === 'light' || theme === 'dark')) {
+    if (
+      (testFunction || (playFunction && args.autoplay !== false)) &&
+      !(theme === 'light' || theme === 'dark')
+    ) {
       theme = 'light';
       showPlayFnNotice = true;
-    } else if (isChromatic() && !storyGlobals.sb_theme && !playFunction) {
+    } else if (isChromatic() && !storyGlobals.sb_theme && !playFunction && !testFunction) {
       theme = 'stacked';
     }
 
@@ -281,8 +284,8 @@ const decorators = [
               <>
                 <PlayFnNotice>
                   <span>
-                    Detected play function in Chromatic. Rendering only light theme to avoid
-                    multiple play functions in the same story.
+                    Detected play/test function in Chromatic. Rendering only light theme to avoid
+                    multiple play/test functions in the same story.
                   </span>
                 </PlayFnNotice>
                 <div style={{ marginBottom: 20 }} />

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

@@ -1,7 +1,7 @@
 import { type RunnerTask, type TaskMeta, type TestContext } from 'vitest';
 
-import { sanitize } from 'storybook/internal/csf';
-import type { ComponentAnnotations, ComposedStoryFn } from 'storybook/internal/types';
+import { type Meta, type Story, isStory, toTestId } from 'storybook/internal/csf';
+import type { ComponentAnnotations, ComposedStoryFn, Renderer } from 'storybook/internal/types';
 
 import { server } from '@vitest/browser/context';
 import { type Report, composeStory, getCsfFactoryAnnotations } from 'storybook/preview-api';
@@ -33,15 +33,19 @@ export const convertToFilePath = (url: string): string => {
 
 export const testStory = (
   exportName: string,
-  story: ComposedStoryFn,
-  meta: ComponentAnnotations,
+  story: ComposedStoryFn | Story<Renderer>,
+  meta: ComponentAnnotations | Meta<Renderer>,
   skipTags: string[],
   testName?: string
 ) => {
   return async (context: TestContext & { story: ComposedStoryFn }) => {
     const annotations = getCsfFactoryAnnotations(story, meta);
+
+    const storyAnnotations =
+      isStory(story) && testName ? story.getAllTests()[testName].story.input : annotations.story;
+
     const composedStory = composeStory(
-      annotations.story,
+      storyAnnotations,
       annotations.meta!,
       { initialGlobals: (await getInitialGlobals?.()) ?? {} },
       annotations.preview ?? globalThis.globalProjectAnnotations,
@@ -59,15 +63,18 @@ export const testStory = (
     };
 
     if (testName) {
-      // TODO: this should be reworked
-      _task.meta.storyId = `${composedStory.id}-${sanitize(testName)}`;
+      _task.meta.storyId = toTestId(composedStory.id, testName);
     } else {
       _task.meta.storyId = composedStory.id;
     }
 
     await setViewport(composedStory.parameters, composedStory.globals);
 
-    await composedStory.run(undefined, testName);
+    if (isStory(story) && testName) {
+      await composedStory.run(undefined, story.getAllTests()[testName].test);
+    } else {
+      await composedStory.run(undefined);
+    }
 
     _task.meta.reports = composedStory.reporting.reports;
   };

code/core/src/component-testing/components/test-fn.stories.tsx

@@ -1,6 +1,6 @@
 import React from 'react';
 
-import type { TestFunction } from 'storybook/internal/types';
+import type { StoryContext } from '@storybook/react-vite';
 
 import { expect, fn } from 'storybook/test';
 
@@ -28,10 +28,13 @@ Default.test('simple', async ({ canvas, userEvent, args }) => {
   await expect(args.onClick).toHaveBeenCalled();
 });
 
-const doTest: TestFunction = async ({ canvas, userEvent, args }) => {
+const doTest = async ({
+  canvas,
+  userEvent,
+  args,
+}: StoryContext<React.ComponentProps<'button'>>) => {
   const button = canvas.getByText('Arg from story');
   await userEvent.click(button);
-  // @ts-expect-error TODO: Fix later with Kasper
   await expect(args.onClick).toHaveBeenCalled();
 };
 Default.test('referring to function in file', doTest);
@@ -40,12 +43,28 @@ Default.test(
   'with overrides',
   {
     args: {
-      children: 'Arg from test',
+      children: 'Arg from test override',
     },
+    parameters: {
+      viewport: {
+        options: {
+          sized: {
+            name: 'Sized',
+            styles: {
+              width: '380px',
+              height: '500px',
+            },
+          },
+        },
+      },
+      chromatic: { viewports: [380] },
+    },
+    globals: { sb_theme: 'dark', viewport: { value: 'sized' } },
   },
   async ({ canvas }) => {
-    const button = canvas.getByText('Arg from test');
+    const button = canvas.getByText('Arg from test override');
     await expect(button).toBeInTheDocument();
+    expect(document.body.clientWidth).toBe(380);
   }
 );
 Default.test(

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

@@ -1,3 +1,4 @@
+//* @vitest-environment happy-dom */
 import { describe, expect, test, vi } from 'vitest';
 
 import { definePreview, definePreviewAddon } from './csf-factories';
@@ -14,7 +15,7 @@ interface Addon2Types {
 
 const addon2 = definePreviewAddon<Addon2Types>({});
 
-const preview = definePreview({ addons: [addon, addon2] });
+const preview = definePreview({ addons: [addon, addon2], renderToCanvas: () => {} });
 
 const meta = preview.meta({
   render: () => 'hello',
@@ -47,18 +48,16 @@ test('addon parameters are inferred', () => {
 describe('test function', () => {
   test('without overrides', async () => {
     const MyStory = meta.story({ args: { label: 'foo' } });
-    const testFn = vi.fn();
+    const testFn = vi.fn(() => console.log('testFn'));
     const testName = 'should run test';
 
     // register test
     MyStory.test(testName, testFn);
-
-    // @ts-expect-error this is a private property not present in the types
-    const storyTest = MyStory.input.__tests![testName];
-    expect(storyTest.input.args).toEqual({ label: 'foo' });
+    const { story: storyTestAnnotations } = MyStory.getAllTests()[testName];
+    expect(storyTestAnnotations.input.args).toEqual({ label: 'foo' });
 
     // execute test
-    await storyTest.input.__testFunction?.();
+    await MyStory.run(undefined, testName);
     expect(testFn).toHaveBeenCalled();
   });
   test('with overrides', async () => {
@@ -68,12 +67,11 @@ describe('test function', () => {
 
     // register test
     MyStory.test(testName, { args: { label: 'bar' } }, testFn);
-    // @ts-expect-error this is a private property not present in the types
-    const storyTest = MyStory.input.__tests![testName];
-    expect(storyTest.input.args).toEqual({ label: 'bar' });
+    const { story: storyTestAnnotations } = MyStory.getAllTests()[testName];
+    expect(storyTestAnnotations.input.args).toEqual({ label: 'bar' });
 
     // execute test
-    await storyTest.input.__testFunction?.();
+    await MyStory.run(undefined, testName);
     expect(testFn).toHaveBeenCalled();
   });
 });

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

@@ -154,6 +154,7 @@ 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> {
@@ -165,7 +166,6 @@ function defineStory<
   TInput extends StoryAnnotations<TRenderer, TRenderer['args']>,
 >(input: TInput, meta: Meta<TRenderer>): Story<TRenderer, TInput> {
   let composed: ComposedStoryFn<TRenderer>;
-  input.__tests ??= {};
   const compose = () => {
     if (!composed) {
       composed = composeStory(
@@ -177,6 +177,9 @@ function defineStory<
     }
     return composed;
   };
+
+  const tests: Record<string, { story: Story<TRenderer>; test: TestFunction<TRenderer> }> = {};
+
   return {
     _tag: 'Story',
     input,
@@ -191,27 +194,22 @@ function defineStory<
     get play() {
       return input.play ?? meta.input?.play ?? (async () => {});
     },
-    get run() {
-      return compose().run ?? (async () => {});
+    async run(context, testName?: string) {
+      const composedRun = compose().run;
+      await composedRun(context, tests[testName!]?.test);
     },
     test(
       name: string,
       overridesOrTestFn: StoryAnnotations<TRenderer, TRenderer['args']> | TestFunction<TRenderer>,
       testFn?: TestFunction<TRenderer>
     ): void {
-      const hasOverrides = typeof overridesOrTestFn !== 'function';
-      const annotations = (hasOverrides ? overridesOrTestFn : {}) as StoryAnnotations<
-        TRenderer,
-        TRenderer['args']
-      >;
-      const testFunction = (hasOverrides ? testFn : overridesOrTestFn) as TestFunction<TRenderer>;
-
-      // A test is a clone of the story + the test function
-      const testStory = this.extend({ ...annotations, tags: ['test-fn'] });
-      testStory.input.__testFunction = testFunction;
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore TODO: Kasper to figure out later
-      this.input.__tests![name] = testStory;
+      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,
+      };
     },
     extend<TInput extends StoryAnnotations<TRenderer, TRenderer['args']>>(input: TInput) {
       return defineStory(
@@ -243,5 +241,8 @@ function defineStory<
         this.meta
       );
     },
+    getAllTests() {
+      return tests;
+    },
   };
 }

code/core/src/csf/story.ts

@@ -538,11 +538,6 @@ export type StoryAnnotations<
 
   /** @deprecated */
   story?: Omit<StoryAnnotations<TRenderer, TArgs>, 'story'>;
-
-  /** @private */
-  // TODO: fix the type issue later
-  __tests?: Record<string, Story<any>>;
-  __testFunction?: TestFunction<TRenderer, TRenderer['args']>;
 } & ({} extends TRequiredArgs ? { args?: TRequiredArgs } : { args: TRequiredArgs });
 
 export type LegacyAnnotatedStoryFn<TRenderer extends Renderer = Renderer, TArgs = Args> = StoryFn<