Don't merge, just testing something out

.gitignore

@@ -21,6 +21,7 @@ test-results
 !/**/.yarn/plugins
 !/**/.yarn/sdks
 !/**/.yarn/versions
+!/**/.yarn/patches
 /**/.pnp.*
 !/node_modules
 

code/addons/actions/src/addArgs.ts

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

code/addons/actions/src/addArgsHelpers.ts

@@ -1,3 +1,4 @@
+/* eslint-disable no-underscore-dangle,no-param-reassign */
 import type { Args, Renderer, ArgsEnhancer } from '@storybook/types';
 import { action } from './runtime/action';
 
@@ -31,7 +32,7 @@ export const inferActionsFromArgTypesRegex: ArgsEnhancer<Renderer> = (context) =
 
   return argTypesMatchingRegex.reduce((acc, [name, argType]) => {
     if (isInInitialArgs(name, initialArgs)) {
-      acc[name] = action(name);
+      acc[name] = action(name, { implicit: true });
     }
     return acc;
   }, {} as Args);
@@ -61,3 +62,33 @@ 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/models/ActionOptions.ts

@@ -4,6 +4,7 @@ interface Options {
   depth: number; // backards compatibility, remove in 7.0
   clearOnStoryChange: boolean;
   limit: number;
+  implicit: boolean;
 }
 
 export type ActionOptions = Partial<Options> & Partial<TelejsonOptions>;

code/addons/actions/src/runtime/action.ts

@@ -46,6 +46,23 @@ export function action(name: string, options: ActionOptions = {}): HandlerFuncti
   };
 
   const handler = function actionHandler(...args: any[]) {
+    // TODO: Enable once codemods are finished
+    // if (options.implicit) {
+    //   const preview =
+    //     '__STORYBOOK_PREVIEW__' in global
+    //       ? (global.__STORYBOOK_PREVIEW__ as PreviewWeb<Renderer>)
+    //       : undefined;
+    //   if (
+    //     preview?.storyRenders.some(
+    //       (render) => render.phase === 'playing' || render.phase === 'rendering'
+    //     )
+    //   ) {
+    //     console.warn(
+    //       'Can not use implicit actions during rendering or playing of a story.'
+    //     );
+    //   }
+    // }
+
     const channel = addons.getChannel();
     const id = uuidv4();
     const minDepth = 5; // anything less is really just storybook internals

code/addons/interactions/src/preview.ts

@@ -1,3 +1,4 @@
+/* eslint-disable no-param-reassign,no-underscore-dangle */
 /// <reference types="node" />
 
 import { addons } from '@storybook/preview-api';
@@ -9,6 +10,7 @@ import type {
   PlayFunction,
   PlayFunctionContext,
   StepLabel,
+  Args,
 } from '@storybook/types';
 import { instrument } from '@storybook/instrumenter';
 import { ModuleMocker } from 'jest-mock';
@@ -33,7 +35,7 @@ const addSpies = (id: string, val: any, key?: string): any => {
   try {
     if (Object.prototype.toString.call(val) === '[object Object]') {
       // We have to mutate the original object for this to survive HMR.
-      // eslint-disable-next-line no-restricted-syntax, no-param-reassign
+      // eslint-disable-next-line no-restricted-syntax
       for (const [k, v] of Object.entries(val)) val[k] = addSpies(id, v, k);
       return val;
     }
@@ -56,7 +58,25 @@ const addSpies = (id: string, val: any, key?: string): any => {
 const addActionsFromArgTypes: ArgsEnhancer<Renderer> = ({ id, initialArgs }) =>
   addSpies(id, initialArgs);
 
-export const argsEnhancers = [addActionsFromArgTypes];
+const instrumentSpies: ArgsEnhancer = ({ initialArgs }) => {
+  const argTypesWithAction = Object.entries(initialArgs).filter(
+    ([, value]) =>
+      typeof value === 'function' &&
+      '_isMockFunction' in value &&
+      value._isMockFunction &&
+      !value._instrumented
+  );
+
+  return argTypesWithAction.reduce((acc, [key, value]) => {
+    const instrumented = instrument({ [key]: () => value }, { retain: true })[key];
+    acc[key] = instrumented();
+    // this enhancer is being called multiple times
+    value._instrumented = true;
+    return acc;
+  }, {} as Args);
+};
+
+export const argsEnhancers = [addActionsFromArgTypes, instrumentSpies];
 
 export const { step: runStep } = instrument(
   {

code/jest.config.base.js

@@ -23,6 +23,7 @@ const modulesToTransform = [
   '@angular',
   '@lit',
   '@mdx-js',
+  '@vitest',
   'ccount',
   'character-entities',
   'decode-named-character-reference',

code/lib/instrumenter/src/instrumenter.test.ts

@@ -112,6 +112,44 @@ describe('Instrumenter', () => {
     expect(result.fn1.fn2.__originalFn__).toBe(fn1.fn2);
   });
 
+  it('patches functions correctly that reference this', () => {
+    const object = {
+      name: 'name',
+      method() {
+        return this.name;
+      },
+    };
+
+    const instrumented = instrument(object);
+    expect(object.method()).toEqual(instrumented.method());
+
+    expect(instrumented.method).toEqual(expect.any(Function));
+    expect(instrumented.method.__originalFn__).toBe(object.method);
+  });
+
+  it('patches functions correctly that use proxies', () => {
+    const object = new Proxy(
+      {
+        name: 'name',
+        method() {
+          return this.name;
+        },
+      },
+      {
+        get(target, prop, receiver) {
+          if (prop === 'name') return `${target[prop]}!`;
+          return Reflect.get(target, prop, receiver);
+        },
+      }
+    );
+
+    const instrumented = instrument(object);
+    expect(object.method()).toEqual(instrumented.method());
+
+    expect(instrumented.method).toEqual(expect.any(Function));
+    expect(instrumented.method.__originalFn__).toBe(object.method);
+  });
+
   it('patched functions call the original function when invoked', () => {
     const { fn } = instrument({ fn: jest.fn() });
     const obj = {};