From 2aff8c5f3a2521b42b7622c83f97299a24982d7d Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Wed, 10 May 2023 15:08:41 +0200
Subject: [PATCH] fix: correctly restore vi.fn implementation (#3341)

---
 packages/spy/src/index.ts        |  2 +-
 test/core/test/jest-mock.test.ts | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts
index 02c9a7e6036..8a14158b4a2 100644
--- 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
index 12bb6fa189a..8d12d16e1ce 100644
--- 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)
+  })
 })