Make sure test packages attaches name to arg spies

storybookjs · Apr 9, 2024 · 726a80a · 726a80a
1 parent 0dc9e58
commit 726a80a
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 54 deletions.
diff --git a/code/addons/interactions/src/preview.ts b/code/addons/interactions/src/preview.ts
@@ -1,11 +1,4 @@
-import type {
-  ArgsEnhancer,
-  PlayFunction,
-  PlayFunctionContext,
-  Renderer,
-  StepLabel,
-} from '@storybook/types';
-import { fn, isMockFunction } from '@storybook/test';
+import type { PlayFunction, PlayFunctionContext, StepLabel } from '@storybook/types';
 import { instrument } from '@storybook/instrumenter';
 
 export const { step: runStep } = instrument(
@@ -16,50 +9,6 @@ export const { step: runStep } = instrument(
   { intercept: true }
 );
 
-export const traverseArgs = (value: unknown, depth = 0, key?: string): unknown => {
-  // Make sure to not get in infinite loops with self referencing args
-  if (depth > 5) return value;
-  if (value == null) return value;
-  if (isMockFunction(value)) {
-    // Makes sure we get the arg name in the interactions panel
-    if (key) value.mockName(key);
-    return value;
-  }
-
-  // wrap explicit actions in a spy
-  if (
-    typeof value === 'function' &&
-    'isAction' in value &&
-    value.isAction &&
-    !('implicit' in value && value.implicit)
-  ) {
-    const mock = fn(value as any);
-    if (key) mock.mockName(key);
-    return mock;
-  }
-
-  if (Array.isArray(value)) {
-    depth++;
-    return value.map((item) => traverseArgs(item, depth));
-  }
-
-  if (typeof value === 'object' && value.constructor === Object) {
-    depth++;
-    for (const [k, v] of Object.entries(value)) {
-      if (Object.getOwnPropertyDescriptor(value, k).writable) {
-        // We have to mutate the original object for this to survive HMR.
-        (value as Record<string, unknown>)[k] = traverseArgs(v, depth, k);
-      }
-    }
-    return value;
-  }
-  return value;
-};
-
-const wrapActionsInSpyFns: ArgsEnhancer<Renderer> = ({ initialArgs }) => traverseArgs(initialArgs);
-
-export const argsEnhancers = [wrapActionsInSpyFns];
-
 export const parameters = {
   throwPlayFunctionExceptions: false,
 };
diff --git a/code/lib/test/src/index.ts b/code/lib/test/src/index.ts
@@ -3,7 +3,15 @@ import { type LoaderFunction } from '@storybook/csf';
 import chai from 'chai';
 import { global } from '@storybook/global';
 import { expect as rawExpect } from './expect';
-import { clearAllMocks, onMockCall, resetAllMocks, restoreAllMocks } from './spy';
+import {
+  clearAllMocks,
+  fn,
+  isMockFunction,
+  onMockCall,
+  resetAllMocks,
+  restoreAllMocks,
+} from './spy';
+import type { Renderer } from '@storybook/types';
 
 export * from './spy';
 
@@ -36,8 +44,52 @@ const resetAllMocksLoader: LoaderFunction = ({ parameters }) => {
   }
 };
 
+export const traverseArgs = (value: unknown, depth = 0, key?: string): unknown => {
+  // Make sure to not get in infinite loops with self referencing args
+  if (depth > 5) return value;
+  if (value == null) return value;
+  if (isMockFunction(value)) {
+    // Makes sure we get the arg name in the interactions panel
+    if (key) value.mockName(key);
+    return value;
+  }
+
+  // wrap explicit actions in a spy
+  if (
+    typeof value === 'function' &&
+    'isAction' in value &&
+    value.isAction &&
+    !('implicit' in value && value.implicit)
+  ) {
+    const mock = fn(value as any);
+    if (key) mock.mockName(key);
+    return mock;
+  }
+
+  if (Array.isArray(value)) {
+    depth++;
+    return value.map((item) => traverseArgs(item, depth));
+  }
+
+  if (typeof value === 'object' && value.constructor === Object) {
+    depth++;
+    for (const [k, v] of Object.entries(value)) {
+      if (Object.getOwnPropertyDescriptor(value, k)?.writable) {
+        // We have to mutate the original object for this to survive HMR.
+        (value as Record<string, unknown>)[k] = traverseArgs(v, depth, k);
+      }
+    }
+    return value;
+  }
+  return value;
+};
+
+const nameSpiesAndWrapActionsInSpies: LoaderFunction<Renderer> = ({ initialArgs }) => {
+  traverseArgs(initialArgs);
+};
+
 // We are using this as a default Storybook loader, when the test package is used. This avoids the need for optional peer dependency workarounds.
 // eslint-disable-next-line no-underscore-dangle
-(global as any).__STORYBOOK_TEST_LOADERS__ = [resetAllMocksLoader];
+(global as any).__STORYBOOK_TEST_LOADERS__ = [resetAllMocksLoader, nameSpiesAndWrapActionsInSpies];
 // eslint-disable-next-line no-underscore-dangle
 (global as any).__STORYBOOK_TEST_ON_MOCK_CALL__ = onMockCall;