Fix mockReset for spies and retain log items originating from the arg enhancer

addons/interactions/src/Panel.tsx

@@ -15,15 +15,23 @@ interface PanelProps {
   active: boolean;
 }
 
-global.window.__STORYBOOK_ADDON_TEST_MANAGER__ = global.window.__STORYBOOK_ADDON_TEST_MANAGER__ || {
-  isDebugging: false,
-  chainedCallIds: new Set<Call['id']>(),
-  playUntil: undefined,
-};
+interface SharedState {
+  isDebugging: boolean;
+  chainedCallIds: Set<Call['id']>;
+  playUntil: Call['id'];
+}
 
-const sharedState = global.window.__STORYBOOK_ADDON_TEST_MANAGER__;
+if (!global.window.__STORYBOOK_ADDON_TEST_MANAGER__) {
+  global.window.__STORYBOOK_ADDON_TEST_MANAGER__ = {
+    isDebugging: false,
+    chainedCallIds: new Set<Call['id']>(),
+    playUntil: undefined,
+  };
+}
 
-const Row = ({
+const sharedState: SharedState = global.window.__STORYBOOK_ADDON_TEST_MANAGER__;
+
+const Interaction = ({
   call,
   callsById,
   onClick,
@@ -116,7 +124,7 @@ const reducer = (
   action: { type: string; payload?: { call?: Call; playUntil?: Call['id'] } }
 ) => {
   switch (action.type) {
-    case 'call':
+    case 'call': {
       const { call } = action.payload;
       const log = state.log.concat(call);
       const interactions = fold(
@@ -134,8 +142,8 @@ const reducer = (
         interactions,
         callsById: { ...state.callsById, [call.id]: call },
       };
-
-    case 'start':
+    }
+    case 'start': {
       sharedState.isDebugging = true;
       sharedState.playUntil = action.payload?.playUntil;
       return {
@@ -144,17 +152,26 @@ const reducer = (
         shadowLog: state.isDebugging ? state.shadowLog : [...state.log],
         isDebugging: true,
       };
-
-    case 'stop':
+    }
+    case 'stop': {
       sharedState.isDebugging = false;
       sharedState.playUntil = undefined;
       return { ...state, isDebugging: false };
-
-    case 'reset':
+    }
+    case 'reset': {
       sharedState.isDebugging = false;
       sharedState.playUntil = undefined;
       sharedState.chainedCallIds.clear();
-      return initialState;
+      const { log, callsById } = state;
+      return {
+        ...initialState,
+        log: log.filter((call) => call.retain),
+        callsById: Object.entries(callsById).reduce<typeof callsById>((acc, [id, call]) => {
+          if (call.retain) acc[id] = call;
+          return acc;
+        }, {}),
+      };
+    }
   }
 };
 
@@ -167,25 +184,27 @@ export const Panel: React.FC<PanelProps> = (props) => {
     },
     setCurrentStory: () => {
       dispatch({ type: 'reset' });
-      emit(EVENTS.RESET);
     },
     storyRendered: () => {
       dispatch({ type: 'stop' });
-      emit(EVENTS.RESET);
     },
   });
 
   const { log, interactions, callsById, isDebugging } = state;
   const hasException = interactions.some((call) => call.state === CallState.ERROR);
-  const hasPrevious = interactions.some(call => call.state !== CallState.PENDING);
+  const hasPrevious = interactions.some((call) => call.state !== CallState.PENDING);
   const hasNext = interactions.some((call) => call.state === CallState.PENDING);
   const nextIndex = interactions.findIndex((call) => call.state === CallState.PENDING);
   const nextCall = interactions[nextIndex];
-  const prevCall = interactions[nextIndex - 2] || (isDebugging ? undefined : interactions.slice(-2)[0]);
+  const prevCall =
+    interactions[nextIndex - 2] || (isDebugging ? undefined : interactions.slice(-2)[0]);
 
   const start = () => {
     const playUntil = log
-      .slice(0, log.findIndex((call) => call.id === interactions[0].id))
+      .slice(
+        0,
+        log.findIndex((call) => call.id === interactions[0].id)
+      )
       .filter((call) => call.interceptable)
       .slice(-1)[0];
     dispatch({ type: 'start', payload: { playUntil: playUntil?.id } });
@@ -200,8 +219,8 @@ export const Panel: React.FC<PanelProps> = (props) => {
       emit(EVENTS.RELOAD);
     }
   };
-  const next = () => goto(nextCall)
-  const prev = () => prevCall ? goto(prevCall) : start()
+  const next = () => goto(nextCall);
+  const prev = () => (prevCall ? goto(prevCall) : start());
   const stop = () => {
     dispatch({ type: 'stop' });
     emit(EVENTS.NEXT);
@@ -222,7 +241,7 @@ export const Panel: React.FC<PanelProps> = (props) => {
       {tabButton && showStatus && ReactDOM.createPortal(statusIcon, tabButton)}
 
       {interactions.map((call) => (
-        <Row call={call} callsById={callsById} key={call.id} onClick={() => goto(call)} />
+        <Interaction call={call} callsById={callsById} key={call.id} onClick={() => goto(call)} />
       ))}
 
       <div style={{ padding: 3 }}>

addons/interactions/src/constants.ts

@@ -7,6 +7,6 @@ export const LOG_STATE_ID = `${ADDON_ID}/state/log`;
 export const EVENTS = {
   CALL: `${ADDON_ID}/call`,
   NEXT: `${ADDON_ID}/next`,
-  RESET: `${ADDON_ID}/reset`,
   RELOAD: `${ADDON_ID}/reload`,
+  SET_CURRENT_STORY: 'setCurrentStory', // Storybook's own event
 };

addons/interactions/src/instrument.test.ts

@@ -21,7 +21,7 @@ global.window = {
 
 beforeEach(() => {
   callSpy.mockReset();
-  addons.getChannel().emit(EVENTS.RESET);
+  addons.getChannel().emit(EVENTS.SET_CURRENT_STORY);
   global.window.__STORYBOOK_ADDON_TEST_PREVIEW__.n = 0;
   global.window.__STORYBOOK_ADDON_TEST_PREVIEW__.next = {};
   global.window.__STORYBOOK_ADDON_TEST_PREVIEW__.callRefsByResult = new Map();
@@ -126,7 +126,7 @@ describe('instrument', () => {
     expect(callSpy).toHaveBeenLastCalledWith(
       expect.objectContaining({
         method: 'fn2',
-        args: [{ __callId__: callSpy.mock.calls[0][0].id }],
+        args: [{ __callId__: callSpy.mock.calls[0][0].id, retain: false }],
       })
     );
   });
@@ -158,10 +158,10 @@ describe('instrument', () => {
           /* call 3 */ 'foo',
           /* call 4 */ 1,
           /* call 5 */ BigInt(1),
-          { __callId__: callSpy.mock.calls[6][0].id },
-          { __callId__: callSpy.mock.calls[7][0].id },
-          { __callId__: callSpy.mock.calls[8][0].id },
-          { __callId__: callSpy.mock.calls[9][0].id },
+          { __callId__: callSpy.mock.calls[6][0].id, retain: false },
+          { __callId__: callSpy.mock.calls[7][0].id, retain: false },
+          { __callId__: callSpy.mock.calls[8][0].id, retain: false },
+          { __callId__: callSpy.mock.calls[9][0].id, retain: false },
         ],
       })
     );
@@ -242,11 +242,11 @@ describe('instrument', () => {
     expect(global.window.location.reload).toHaveBeenCalled();
   });
 
-  it('resets preview state on the "reset" event', () => {
+  it('resets preview state when switching stories', () => {
     global.window.__STORYBOOK_ADDON_TEST_PREVIEW__.n = 123;
     global.window.__STORYBOOK_ADDON_TEST_PREVIEW__.next = { ref: () => {} };
     global.window.__STORYBOOK_ADDON_TEST_PREVIEW__.callRefsByResult = new Map([[{}, 'ref']]);
-    addons.getChannel().emit(EVENTS.RESET);
+    addons.getChannel().emit(EVENTS.SET_CURRENT_STORY);
     expect(global.window.__STORYBOOK_ADDON_TEST_PREVIEW__).toStrictEqual({
       n: 0,
       next: {},

addons/interactions/src/instrument.ts

@@ -7,6 +7,7 @@ import { Call, CallRef, CallState } from './types';
 
 export interface Options {
   intercept?: boolean;
+  retain?: boolean;
   mutate?: boolean;
   path?: Array<string | CallRef>;
 }
@@ -43,10 +44,12 @@ const init = () => {
 
   channel.on(EVENTS.NEXT, () => Object.values(iframeState.next).forEach((resolve) => resolve()));
   channel.on(EVENTS.RELOAD, () => global.window.location.reload());
-  channel.on(EVENTS.RESET, () => {
-    iframeState.n = 0;
+  channel.on(EVENTS.SET_CURRENT_STORY, () => {
+    iframeState.callRefsByResult = new Map(
+      Array.from(iframeState.callRefsByResult.entries()).filter(([, val]) => val.retain)
+    );
+    iframeState.n = iframeState.callRefsByResult.size;
     iframeState.next = {};
-    iframeState.callRefsByResult.clear();
   });
 };
 
@@ -95,7 +98,7 @@ function invoke(fn: Function, call: Call) {
     if (result && ['object', 'function', 'symbol'].includes(typeof result)) {
       // Track the result so we can trace later uses of it back to the originating call.
       // Primitive results (undefined, null, boolean, string, number, BigInt) are ignored.
-      iframeState.callRefsByResult.set(result, { __callId__: call.id });
+      iframeState.callRefsByResult.set(result, { __callId__: call.id, retain: call.retain });
     }
     return result;
   } catch (e) {
@@ -129,7 +132,8 @@ function intercept(fn: Function, call: Call) {
 // returns the original result.
 function track(method: string, fn: Function, args: any[], { path = [], ...options }: Options) {
   const id = `${iframeState.n++}-${method}`;
-  const call: Call = { id, path, method, args, interceptable: !!options.intercept };
+  const { intercept: interceptable = false, retain = false } = options;
+  const call: Call = { id, path, method, args, interceptable, retain };
   const result = (options.intercept ? intercept : invoke)(fn, call);
   return instrument(result, { ...options, mutate: true, path: [{ __callId__: call.id }] });
 }

addons/interactions/src/preset/preview.ts

@@ -1,19 +1,20 @@
 import { Args, addons } from '@storybook/addons';
 import { ArgsEnhancer } from '@storybook/client-api';
-import { jest } from '../jest-mock';
+import { fn } from 'jest-mock';
 import { EVENTS } from '../constants';
+import { instrument } from '../instrument';
 
+const { fn: spy } = instrument({ fn }, { retain: true, path: ['jest'] });
 const channel = addons.getChannel();
+const spies: any[] = [];
 
-const addActionsFromArgTypes: ArgsEnhancer = ({ args }) => {
-  const mocks: any[] = [];
-  channel.on(EVENTS.RESET, () => {
-    // mocks.forEach(mock => mock.mockReset())
-  });
+channel.on(EVENTS.SET_CURRENT_STORY, () => spies.forEach((mock) => mock.mockReset()));
+
+const addActionsFromArgTypes: ArgsEnhancer = ({ id, args }) => {
   return Object.entries(args).reduce((acc, [key, val]) => {
     if (typeof val === 'function' && val.name === 'actionHandler') {
-      acc[key] = jest.fn(val);
-      mocks.push(acc[key]);
+      acc[key] = spy(val);
+      spies.push(acc[key]);
       return acc;
     }
     acc[key] = val;

addons/interactions/src/types.ts

@@ -4,12 +4,14 @@ export interface Call {
   method: string;
   args: any[];
   interceptable: boolean;
+  retain: boolean;
   state?: `${CallState}`;
   exception?: CaughtException;
 }
 
 export interface CallRef {
   __callId__: Call['id'];
+  retain?: boolean;
 }
 
 export enum CallState {