From 2cb1a15a5874ed07396a09c4180ed1911d6b1b00 Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Wed, 17 May 2023 15:06:55 +0200
Subject: [PATCH] fix: stop spying on a method, when it's restored (#3386)

---
 packages/spy/src/index.ts        | 17 +++++++++++------
 test/core/test/jest-mock.test.ts |  2 +-
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts
index 8a14158b4a2..d077893df3e 100644
--- a/packages/spy/src/index.ts
+++ b/packages/spy/src/index.ts
@@ -213,6 +213,13 @@ function enhanceSpy<TArgs extends any[], TReturns>(
   let onceImplementations: ((...args: TArgs) => TReturns)[] = []
   let implementationChangedTemporarily = false
 
+  function mockCall(this: unknown, ...args: any) {
+    instances.push(this)
+    invocations.push(++callOrder)
+    const impl = implementationChangedTemporarily ? implementation! : (onceImplementations.shift() || implementation || state.getOriginal() || (() => {}))
+    return impl.apply(this, args)
+  }
+
   let name: string = (stub as any).name
 
   stub.getMockName = () => name || 'vi.fn()'
@@ -237,6 +244,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
 
   stub.mockRestore = () => {
     stub.mockReset()
+    state.restore()
     implementation = undefined
     return stub
   }
@@ -244,6 +252,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
   stub.getMockImplementation = () => implementation
   stub.mockImplementation = (fn: (...args: TArgs) => TReturns) => {
     implementation = fn
+    state.willCall(mockCall)
     return stub
   }
 
@@ -258,6 +267,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
     const originalImplementation = implementation
 
     implementation = fn
+    state.willCall(mockCall)
     implementationChangedTemporarily = true
 
     const reset = () => {
@@ -305,12 +315,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
     get: () => mockContext,
   })
 
-  state.willCall(function (this: unknown, ...args) {
-    instances.push(this)
-    invocations.push(++callOrder)
-    const impl = implementationChangedTemporarily ? implementation! : (onceImplementations.shift() || implementation || state.getOriginal() || (() => {}))
-    return impl.apply(this, args)
-  })
+  state.willCall(mockCall)
 
   spies.add(stub)
 
diff --git a/test/core/test/jest-mock.test.ts b/test/core/test/jest-mock.test.ts
index 8d12d16e1ce..e460a1a4295 100644
--- a/test/core/test/jest-mock.test.ts
+++ b/test/core/test/jest-mock.test.ts
@@ -261,7 +261,7 @@ describe('jest mock compat layer', () => {
     obj.property = true
     // unlike jest, mockRestore only restores implementation to the original one,
     // we are still spying on the setter
-    expect(spy).toHaveBeenCalled()
+    expect(spy).not.toHaveBeenCalled()
     expect(obj.property).toBe(true)
   })