fix: clear mock from callstack, if mocked with factory (#1276)

* fix: clear mock from callstack, if mocked with factory

* test: fix test for zustand

examples/mocks/src/zustand.ts → examples/mocks/__mocks__/zustand.ts

@@ -1,13 +1,8 @@
 import actualCreate from 'zustand'
 
-// a variable to hold reset functions for all stores declared in the app
-const storeResetFns = new Set()
-
 // when creating a store, we get its initial state, create a reset function and add it in the set
 const create = vi.fn((createState) => {
   const store = actualCreate(createState)
-  const initialState = store.getState()
-  storeResetFns.add(() => store.setState(initialState, true))
   return store
 })
 

examples/mocks/src/zustand-magic.ts

@@ -0,0 +1,5 @@
+import zustand from 'zustand'
+
+export const magic = () => {
+  return zustand()
+}

examples/mocks/test/self-importing.test.ts

@@ -1,9 +1,16 @@
 import zustand from 'zustand'
+import { magic } from '../src/zustand-magic'
 
 vi.mock('zustand')
 
 describe('zustand didn\'t go into an infinite loop', () => {
   test('zustand is mocked', () => {
     expect(vi.isMockFunction(zustand)).toBe(true)
   })
+
+  test('magic calls zustand', () => {
+    const store = magic()
+    expect(zustand).toHaveBeenCalled()
+    expect(store).toBeTypeOf('function')
+  })
 })

packages/vitest/src/runtime/mocker.ts

@@ -250,7 +250,10 @@ export class VitestMocker {
     }
     if (typeof mock === 'function' && !callstack.includes(`mock:${dep}`)) {
       callstack.push(`mock:${dep}`)
-      return this.callFunctionMock(dep, mock)
+      const result = await this.callFunctionMock(dep, mock)
+      const indexMock = callstack.indexOf(`mock:${dep}`)
+      callstack.splice(indexMock, 1)
+      return result
     }
     if (typeof mock === 'string' && !callstack.includes(mock))
       dep = mock