fix: correctly restore vi.fn implementation (#3341)

vitest-dev · May 10, 2023 · 2aff8c5 · 2aff8c5
1 parent ddbba39
commit 2aff8c5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts
@@ -324,7 +324,7 @@ export function fn<TArgs extends any[] = any[], R = any>(
 export function fn<TArgs extends any[] = any[], R = any>(
   implementation?: (...args: TArgs) => R,
 ): Mock<TArgs, R> {
-  const enhancedSpy = enhanceSpy(tinyspy.internalSpyOn({ spy: () => {} }, 'spy'))
+  const enhancedSpy = enhanceSpy(tinyspy.internalSpyOn({ spy: implementation || (() => {}) }, 'spy'))
   if (implementation)
     enhancedSpy.mockImplementation(implementation)
 

diff --git a/test/core/test/jest-mock.test.ts b/test/core/test/jest-mock.test.ts
@@ -321,4 +321,15 @@ describe('jest mock compat layer', () => {
     const instance2 = new Fn()
     expect(Fn.mock.instances[1]).toBe(instance2)
   })
+
+  it('.mockRestore() should restore initial implementation', () => {
+    const testFn = vi.fn(() => true)
+    expect(testFn()).toBe(true)
+
+    testFn.mockReturnValue(false)
+    expect(testFn()).toBe(false)
+
+    testFn.mockRestore()
+    expect(testFn()).toBe(true)
+  })
 })