Test: Model loaders as before each and restore mocks properly

code/addons/actions/src/addArgs.ts

@@ -1,12 +1,7 @@
 import type { ArgsEnhancer } from '@storybook/types';
-import {
-  addActionsFromArgTypes,
-  attachActionsToFunctionMocks,
-  inferActionsFromArgTypesRegex,
-} from './addArgsHelpers';
+import { addActionsFromArgTypes, inferActionsFromArgTypesRegex } from './addArgsHelpers';
 
 export const argsEnhancers: ArgsEnhancer[] = [
   addActionsFromArgTypes,
   inferActionsFromArgTypesRegex,
-  attachActionsToFunctionMocks,
 ];

code/addons/actions/src/addArgsHelpers.ts

@@ -1,4 +1,3 @@
-/* eslint-disable no-underscore-dangle,no-param-reassign */
 import type { Args, Renderer, ArgsEnhancer } from '@storybook/types';
 import { action } from './runtime/action';
 
@@ -63,33 +62,3 @@ export const addActionsFromArgTypes: ArgsEnhancer<Renderer> = (context) => {
     return acc;
   }, {} as Args);
 };
-
-export const attachActionsToFunctionMocks: ArgsEnhancer<Renderer> = (context) => {
-  const {
-    initialArgs,
-    argTypes,
-    parameters: { actions },
-  } = context;
-  if (actions?.disable || !argTypes) {
-    return {};
-  }
-
-  const argTypesWithAction = Object.entries(initialArgs).filter(
-    ([, value]) =>
-      typeof value === 'function' &&
-      '_isMockFunction' in value &&
-      value._isMockFunction &&
-      !value._actionAttached
-  );
-
-  return argTypesWithAction.reduce((acc, [key, value]) => {
-    const previous = value.getMockImplementation();
-    value.mockImplementation((...args: unknown[]) => {
-      action(key)(...args);
-      return previous?.(...args);
-    });
-    // this enhancer is being called multiple times
-    value._actionAttached = true;
-    return acc;
-  }, {} as Args);
-};

code/addons/actions/src/loaders.ts

@@ -0,0 +1,30 @@
+/* eslint-disable no-underscore-dangle */
+import type { LoaderFunction } from '@storybook/types';
+import { action } from './runtime';
+
+const attachActionsToFunctionMocks: LoaderFunction = (context) => {
+  const {
+    args,
+    parameters: { actions },
+  } = context;
+  if (actions?.disable) return;
+
+  Object.entries(args)
+    .filter(
+      ([, value]) =>
+        typeof value === 'function' && '_isMockFunction' in value && value._isMockFunction
+    )
+    .forEach(([key, value]) => {
+      const previous = value.getMockImplementation();
+      if (previous?._actionAttached !== true) {
+        const implementation = (...params: unknown[]) => {
+          action(key)(...params);
+          return previous?.(...params);
+        };
+        implementation._actionAttached = true;
+        value.mockImplementation(implementation);
+      }
+    });
+};
+
+export const loaders: LoaderFunction[] = [attachActionsToFunctionMocks];

code/addons/actions/src/preview.ts

@@ -1 +1,2 @@
 export * from './addArgs';
+export * from './loaders';

code/addons/links/package.json

@@ -63,7 +63,7 @@
     "prep": "node --loader ../../../scripts/node_modules/esbuild-register/loader.js -r ../../../scripts/node_modules/esbuild-register/register.js ../../../scripts/prepare/addon-bundle.ts"
   },
   "dependencies": {
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/global": "^5.0.0",
     "ts-dedent": "^2.0.0"
   },

code/addons/storyshots-puppeteer/package.json

@@ -37,7 +37,7 @@
   },
   "dependencies": {
     "@axe-core/puppeteer": "^4.2.0",
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/node-logger": "workspace:*",
     "@storybook/types": "workspace:*",
     "@types/jest-image-snapshot": "^6.0.0",

code/lib/codemod/package.json

@@ -55,7 +55,7 @@
     "@babel/core": "^7.23.2",
     "@babel/preset-env": "^7.23.2",
     "@babel/types": "^7.23.0",
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/csf-tools": "workspace:*",
     "@storybook/node-logger": "workspace:*",
     "@storybook/types": "workspace:*",

code/lib/core-server/package.json

@@ -66,7 +66,7 @@
     "@storybook/channels": "workspace:*",
     "@storybook/core-common": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/csf-tools": "workspace:*",
     "@storybook/docs-mdx": "^0.1.0",
     "@storybook/global": "^5.0.0",

code/lib/csf-tools/package.json

@@ -46,7 +46,7 @@
     "@babel/parser": "^7.23.0",
     "@babel/traverse": "^7.23.2",
     "@babel/types": "^7.23.0",
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/types": "workspace:*",
     "fs-extra": "^11.1.0",
     "recast": "^0.23.1",

code/lib/manager-api/package.json

@@ -46,7 +46,7 @@
     "@storybook/channels": "workspace:*",
     "@storybook/client-logger": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/global": "^5.0.0",
     "@storybook/router": "workspace:*",
     "@storybook/theming": "workspace:*",

code/lib/preview-api/package.json

@@ -71,7 +71,7 @@
     "@storybook/channels": "workspace:*",
     "@storybook/client-logger": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "^0.1.0",
+    "@storybook/csf": "^0.1.2",
     "@storybook/global": "^5.0.0",
     "@storybook/types": "workspace:*",
     "@types/qs": "^6.9.5",

code/lib/preview-api/src/modules/client-api/ClientApi.ts

@@ -168,9 +168,11 @@ export class ClientApi<TRenderer extends Renderer> {
     }
   };
 
-  addStepRunner = (stepRunner: StepRunner) => {
+  addStepRunner = (stepRunner: StepRunner<TRenderer>) => {
     this.facade.projectAnnotations.runStep = composeStepRunners(
-      [this.facade.projectAnnotations.runStep, stepRunner].filter(Boolean) as StepRunner[]
+      [this.facade.projectAnnotations.runStep, stepRunner].filter(
+        Boolean
+      ) as StepRunner<TRenderer>[]
     );
   };
 
@@ -297,7 +299,7 @@ export class ClientApi<TRenderer extends Renderer> {
     this._addedExports[fileName] = { default: meta };
 
     let counter = 0;
-    api.add = (storyName: string, storyFn: StoryFn<TRenderer>, parameters: Parameters = {}) => {
+    api.add = (storyName: string, storyFn: StoryFn, parameters: Parameters = {}) => {
       hasAdded = true;
 
       if (typeof storyName !== 'string') {

code/lib/preview-api/src/modules/preview-web/docs-context/DocsContext.ts

@@ -34,7 +34,7 @@ export class DocsContext<TRenderer extends Renderer> implements DocsContextProps
   constructor(
     public channel: Channel,
     protected store: StoryStore<TRenderer>,
-    public renderStoryToElement: DocsContextProps['renderStoryToElement'],
+    public renderStoryToElement: DocsContextProps<TRenderer>['renderStoryToElement'],
     /** The CSF files known (via the index) to be refererenced by this docs file */
     csfFiles: CSFFile<TRenderer>[]
   ) {

code/lib/preview-api/src/modules/preview-web/render/CsfDocsRender.ts

@@ -95,7 +95,7 @@ export class CsfDocsRender<TRenderer extends Renderer> implements Render<TRender
     );
   }
 
-  docsContext(renderStoryToElement: DocsContextProps['renderStoryToElement']) {
+  docsContext(renderStoryToElement: DocsContextProps<TRenderer>['renderStoryToElement']) {
     if (!this.csfFiles) throw new Error('Cannot render docs before preparing');
     const docsContext = new DocsContext<TRenderer>(
       this.channel,
@@ -112,7 +112,7 @@ export class CsfDocsRender<TRenderer extends Renderer> implements Render<TRender
 
   async renderToElement(
     canvasElement: TRenderer['canvasElement'],
-    renderStoryToElement: DocsContextProps['renderStoryToElement']
+    renderStoryToElement: DocsContextProps<TRenderer>['renderStoryToElement']
   ) {
     if (!this.story || !this.csfFiles) throw new Error('Cannot render docs before preparing');
 

code/lib/preview-api/src/modules/preview-web/render/MdxDocsRender.ts

@@ -79,7 +79,7 @@ export class MdxDocsRender<TRenderer extends Renderer> implements Render<TRender
     );
   }
 
-  docsContext(renderStoryToElement: DocsContextProps['renderStoryToElement']) {
+  docsContext(renderStoryToElement: DocsContextProps<TRenderer>['renderStoryToElement']) {
     if (!this.csfFiles) throw new Error('Cannot render docs before preparing');
 
     // NOTE we do *not* attach any CSF file yet. We wait for `referenceMeta(..., true)`
@@ -94,7 +94,7 @@ export class MdxDocsRender<TRenderer extends Renderer> implements Render<TRender
 
   async renderToElement(
     canvasElement: TRenderer['canvasElement'],
-    renderStoryToElement: DocsContextProps['renderStoryToElement']
+    renderStoryToElement: DocsContextProps<TRenderer>['renderStoryToElement']
   ) {
     if (!this.exports || !this.csfFiles || !this.store.projectAnnotations)
       throw new Error('Cannot render docs before preparing');

code/lib/preview-api/src/modules/store/csf/normalizeArrays.ts

@@ -0,0 +1,4 @@
+export const normalizeArrays = <T>(array: T[] | T | undefined): T[] => {
+  if (Array.isArray(array)) return array;
+  return array ? [array] : [];
+};

code/lib/preview-api/src/modules/store/csf/normalizeProjectAnnotations.ts

@@ -8,16 +8,21 @@ import type {
 import { inferArgTypes } from '../inferArgTypes';
 import { inferControls } from '../inferControls';
 import { normalizeInputTypes } from './normalizeInputTypes';
+import { normalizeArrays } from './normalizeArrays';
 
 export function normalizeProjectAnnotations<TRenderer extends Renderer>({
   argTypes,
   globalTypes,
   argTypesEnhancers,
+  decorators,
+  loaders,
   ...annotations
 }: ProjectAnnotations<TRenderer>): NormalizedProjectAnnotations<TRenderer> {
   return {
     ...(argTypes && { argTypes: normalizeInputTypes(argTypes as ArgTypes) }),
     ...(globalTypes && { globalTypes: normalizeInputTypes(globalTypes) }),
+    decorators: normalizeArrays(decorators),
+    loaders: normalizeArrays(loaders),
     argTypesEnhancers: [
       ...(argTypesEnhancers || []),
       inferArgTypes,

code/lib/preview-api/src/modules/store/csf/normalizeStory.ts

@@ -13,6 +13,7 @@ import { dedent } from 'ts-dedent';
 import { logger } from '@storybook/client-logger';
 import deprecate from 'util-deprecate';
 import { normalizeInputTypes } from './normalizeInputTypes';
+import { normalizeArrays } from './normalizeArrays';
 
 const deprecatedStoryAnnotation = dedent`
 CSF .story annotations deprecated; annotate story functions directly:
@@ -44,11 +45,15 @@ export function normalizeStory<TRenderer extends Renderer>(
     storyObject.storyName ||
     story?.name ||
     exportName;
-  const decorators = [...(storyObject.decorators || []), ...(story?.decorators || [])];
+
+  const decorators = [
+    ...normalizeArrays(storyObject.decorators),
+    ...normalizeArrays(story?.decorators),
+  ];
   const parameters = { ...story?.parameters, ...storyObject.parameters };
   const args = { ...story?.args, ...storyObject.args };
   const argTypes = { ...(story?.argTypes as ArgTypes), ...(storyObject.argTypes as ArgTypes) };
-  const loaders = [...(storyObject.loaders || []), ...(story?.loaders || [])];
+  const loaders = [...normalizeArrays(storyObject.loaders), ...normalizeArrays(story?.loaders)];
   const { render, play, tags = [] } = storyObject;
 
   // eslint-disable-next-line no-underscore-dangle

code/lib/preview-api/src/modules/store/csf/prepareStory.test.ts

@@ -343,9 +343,9 @@ describe('prepareStory', () => {
       );
 
       const storyContext = { context: 'value' } as any;
-      const loadedContext = await applyLoaders(storyContext);
+      const loadedContext = await applyLoaders({ ...storyContext });
 
-      expect(loader).toHaveBeenCalledWith(storyContext);
+      expect(loader).toHaveBeenCalledWith({ ...storyContext, loaded: {} });
       expect(loadedContext).toEqual({
         context: 'value',
         loaded: { foo: 7 },

code/lib/preview-api/src/modules/store/csf/prepareStory.ts

@@ -1,31 +1,33 @@
+/* eslint-disable no-restricted-syntax,no-await-in-loop,@typescript-eslint/no-loop-func,no-underscore-dangle */
 import { global } from '@storybook/global';
 
 import type {
-  Renderer,
   Args,
   ArgsStoryFn,
   LegacyStoryFn,
-  Parameters,
-  PlayFunction,
-  PlayFunctionContext,
-  StepLabel,
+  ModuleExport,
   NormalizedComponentAnnotations,
   NormalizedProjectAnnotations,
   NormalizedStoryAnnotations,
+  Parameters,
+  PlayFunction,
+  PlayFunctionContext,
+  PreparedMeta,
   PreparedStory,
+  Renderer,
+  StepLabel,
   StoryContext,
   StoryContextForEnhancers,
   StoryContextForLoaders,
   StrictArgTypes,
-  PreparedMeta,
-  ModuleExport,
 } from '@storybook/types';
 import { includeConditionalArg } from '@storybook/csf';
 
 import { applyHooks } from '../../addons';
 import { combineParameters } from '../parameters';
 import { defaultDecorateStory } from '../decorators';
 import { groupArgsByTarget, UNTARGETED } from '../args';
+import { normalizeArrays } from './normalizeArrays';
 
 // Combine all the metadata about a story (both direct and inherited from the component/global scope)
 // into a "renderable" story function, with all decorators applied, parameters passed as context etc
@@ -48,15 +50,23 @@ export function prepareStory<TRenderer extends Renderer>(
     projectAnnotations
   );
 
-  const loaders = [
-    ...(projectAnnotations.loaders || []),
-    ...(componentAnnotations.loaders || []),
-    ...(storyAnnotations?.loaders || []),
-  ];
-  const applyLoaders = async (context: StoryContextForLoaders<TRenderer>) => {
-    const loadResults = await Promise.all(loaders.map((loader) => loader(context)));
-    const loaded = Object.assign({}, ...loadResults);
-    return { ...context, loaded };
+  const applyLoaders = async (
+    context: StoryContextForLoaders<TRenderer>
+  ): Promise<StoryContextForLoaders<TRenderer> & { loaded: StoryContext<TRenderer>['loaded'] }> => {
+    let updatedContext = { ...context, loaded: {} };
+    for (const loaders of [
+      ...('__STORYBOOK_TEST_LOADERS__' in global && Array.isArray(global.__STORYBOOK_TEST_LOADERS__)
+        ? [global.__STORYBOOK_TEST_LOADERS__]
+        : []),
+      normalizeArrays(projectAnnotations.loaders),
+      normalizeArrays(componentAnnotations.loaders),
+      normalizeArrays(storyAnnotations.loaders),
+    ]) {
+      const loadResults = await Promise.all(loaders.map((loader) => loader(updatedContext)));
+      const loaded: Record<string, any> = Object.assign({}, ...loadResults);
+      updatedContext = { ...updatedContext, loaded: { ...updatedContext.loaded, ...loaded } };
+    }
+    return updatedContext;
   };
 
   const undecoratedStoryFn: LegacyStoryFn<TRenderer> = (context: StoryContext<TRenderer>) => {
@@ -70,9 +80,9 @@ export function prepareStory<TRenderer extends Renderer>(
   const { applyDecorators = defaultDecorateStory, runStep } = projectAnnotations;
 
   const decorators = [
-    ...(storyAnnotations?.decorators || []),
-    ...(componentAnnotations.decorators || []),
-    ...(projectAnnotations.decorators || []),
+    ...normalizeArrays(storyAnnotations?.decorators),
+    ...normalizeArrays(componentAnnotations?.decorators),
+    ...normalizeArrays(projectAnnotations?.decorators),
   ];
 
   // The render function on annotations *has* to be an `ArgsStoryFn`, so when we normalize
@@ -115,7 +125,6 @@ export function prepareStory<TRenderer extends Renderer>(
     playFunction,
   };
 }
-
 export function prepareMeta<TRenderer extends Renderer>(
   componentAnnotations: NormalizedComponentAnnotations<TRenderer>,
   projectAnnotations: NormalizedProjectAnnotations<TRenderer>,
@@ -164,7 +173,6 @@ function preparePartialAnnotations<TRenderer extends Renderer>(
 
     const { passArgsFirst = true } = parameters;
 
-    // eslint-disable-next-line no-underscore-dangle
     parameters.__isArgsStory = passArgsFirst && render && render.length > 0;
   }
 

code/lib/preview-api/src/modules/store/csf/testing-utils/index.ts

@@ -82,7 +82,7 @@ export function composeStory<TRenderer extends Renderer = Renderer, TArgs extend
         args: { ...story.initialArgs, ...extraArgs },
       };
 
-      return story.unboundStoryFn(prepareContext(context as StoryContext));
+      return story.unboundStoryFn(prepareContext(context as StoryContext<TRenderer>));
     },
     {
       storyName,