Move args enhancer to be a loader

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';