Test: Make spies reactive so that they can be logged by addon-actions

code/addons/actions/src/loaders.ts

@@ -1,36 +1,26 @@
 /* eslint-disable no-underscore-dangle */
 import type { LoaderFunction } from '@storybook/types';
+import { global } from '@storybook/global';
+import type { onMockCall as onMockCallType } from '@storybook/test';
 import { action } from './runtime';
 
-export const tinySpyInternalState = Symbol.for('tinyspy:spy');
+let subscribed = false;
 
-const attachActionsToFunctionMocks: LoaderFunction = (context) => {
+const logActionsWhenMockCalled: 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]) => {
-      // See this discussion for context:
-      // https://github.com/vitest-dev/vitest/pull/5352
-      const previous =
-        value.getMockImplementation() ??
-        (tinySpyInternalState in value ? value[tinySpyInternalState]?.getOriginal() : undefined);
-      if (previous?._actionAttached !== true && previous?.isAction !== true) {
-        const implementation = (...params: unknown[]) => {
-          action(key)(...params);
-          return previous?.(...params);
-        };
-        implementation._actionAttached = true;
-        value.mockImplementation(implementation);
-      }
-    });
+  if (
+    !subscribed &&
+    '__STORYBOOK_TEST_ON_MOCK_CALL__' in global &&
+    typeof global.__STORYBOOK_TEST_ON_MOCK_CALL__ === 'function'
+  ) {
+    const onMockCall = global.__STORYBOOK_TEST_ON_MOCK_CALL__ as typeof onMockCallType;
+    onMockCall((mock, args) => action(mock.getMockName())(args));
+    subscribed = true;
+  }
 };
 
-export const loaders: LoaderFunction[] = [attachActionsToFunctionMocks];
+export const loaders: LoaderFunction[] = [logActionsWhenMockCalled];

code/lib/test/package.json

@@ -56,6 +56,7 @@
     "util": "^0.12.4"
   },
   "devDependencies": {
+    "tinyspy": "^2.2.0",
     "ts-dedent": "^2.2.0",
     "type-fest": "~2.19",
     "typescript": "^5.3.2"

code/lib/test/src/index.ts

@@ -3,7 +3,7 @@ import { type LoaderFunction } from '@storybook/csf';
 import chai from 'chai';
 import { global } from '@storybook/global';
 import { expect as rawExpect } from './expect';
-import { clearAllMocks, resetAllMocks, restoreAllMocks } from './spy';
+import { clearAllMocks, onMockCall, resetAllMocks, restoreAllMocks } from './spy';
 
 export * from './spy';
 
@@ -39,3 +39,5 @@ const resetAllMocksLoader: LoaderFunction = ({ parameters }) => {
 // 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];
+// eslint-disable-next-line no-underscore-dangle
+(global as any).__STORYBOOK_TEST_ON_MOCK_CALL__ = onMockCall;

code/lib/test/src/spy.test.ts

@@ -0,0 +1,15 @@
+import { it, vi, expect, beforeEach } from 'vitest';
+import { fn, onMockCall } from './spy';
+
+const vitestSpy = vi.fn();
+
+beforeEach(() => {
+  const unsubscribe = onMockCall(vitestSpy);
+  return () => unsubscribe();
+});
+
+it('mocks are reactive', () => {
+  const storybookSpy = fn();
+  storybookSpy(1);
+  expect(vitestSpy).toHaveBeenCalledWith(storybookSpy, [1]);
+});

code/lib/test/src/spy.ts

@@ -1,17 +1,63 @@
+/* eslint-disable @typescript-eslint/no-shadow */
+import type { Mock, MockInstance } from '@vitest/spy';
 import {
-  spyOn,
+  spyOn as vitestSpyOn,
   isMockFunction,
-  fn,
+  fn as vitestFn,
   mocks,
   type MaybeMocked,
   type MaybeMockedDeep,
   type MaybePartiallyMocked,
   type MaybePartiallyMockedDeep,
 } from '@vitest/spy';
+import type { SpyInternalImpl } from 'tinyspy';
+import * as tinyspy from 'tinyspy';
 
 export type * from '@vitest/spy';
 
-export { spyOn, isMockFunction, fn, mocks };
+export { isMockFunction, mocks };
+
+type Listener = (mock: MockInstance, args: unknown[]) => void;
+let listeners: Listener[] = [];
+
+export function onMockCall(callback: Listener): () => void {
+  listeners = [...listeners, callback];
+  return () => {
+    listeners = listeners.filter((listener) => listener !== callback);
+  };
+}
+
+// @ts-expect-error TS is hard you know
+export const spyOn: typeof vitestSpyOn = (...args) => {
+  const mock = vitestSpyOn(...(args as Parameters<typeof vitestSpyOn>));
+  return reactiveMock(mock);
+};
+
+export function fn<TArgs extends any[] = any, R = any>(): Mock<TArgs, R>;
+export function fn<TArgs extends any[] = any[], R = any>(
+  implementation: (...args: TArgs) => R
+): Mock<TArgs, R>;
+export function fn<TArgs extends any[] = any[], R = any>(implementation?: (...args: TArgs) => R) {
+  const mock = implementation ? vitestFn(implementation) : vitestFn();
+  return reactiveMock(mock);
+}
+
+function reactiveMock(mock: MockInstance) {
+  const reactive = listenWhenCalled(mock);
+  const originalMockImplementation = reactive.mockImplementation.bind(null);
+  reactive.mockImplementation = (fn) => listenWhenCalled(originalMockImplementation(fn));
+  return reactive;
+}
+
+function listenWhenCalled(mock: MockInstance) {
+  const state = tinyspy.getInternalState(mock as unknown as SpyInternalImpl);
+  const impl = state.impl?.bind(null);
+  state.willCall((...args) => {
+    listeners.forEach((listener) => listener(mock, args));
+    return impl?.(...args);
+  });
+  return mock;
+}
 
 /**
  * Calls [`.mockClear()`](https://vitest.dev/api/mock#mockclear) on every mocked function. This will only empty `.mock` state, it will not reset implementation.

code/yarn.lock

@@ -6667,6 +6667,7 @@ __metadata:
     "@vitest/expect": "npm:1.3.1"
     "@vitest/spy": "npm:^1.3.1"
     chai: "npm:^4.4.1"
+    tinyspy: "npm:^2.2.0"
     ts-dedent: "npm:^2.2.0"
     type-fest: "npm:~2.19"
     typescript: "npm:^5.3.2"