fix: stop spying on a method, when it's restored (#3386)

vitest-dev · May 17, 2023 · 2cb1a15 · 2cb1a15
1 parent dcf1346
commit 2cb1a15
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git 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,13 +244,15 @@ function enhanceSpy<TArgs extends any[], TReturns>(
 
   stub.mockRestore = () => {
     stub.mockReset()
+    state.restore()
     implementation = undefined
     return stub
   }
 
   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
@@ -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)
   })