From b3e43d0452ac463e4fa521d6e6b5480bea097274 Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Mon, 9 Dec 2024 10:24:20 +0100
Subject: [PATCH] fix(spy)!: spyOn reuses mock if method is already spyed on
 (#6464)

Co-authored-by: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
---
 packages/mocker/src/automocker.ts |  2 +-
 packages/spy/src/index.ts         | 25 +++++++++++++++++++++++++
 test/core/test/jest-mock.test.ts  | 25 +++++++++++++++++++++++++
 test/core/test/mocked.test.ts     |  3 +++
 4 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/packages/mocker/src/automocker.ts b/packages/mocker/src/automocker.ts
index 571079e107e0..256dbb952310 100644
--- a/packages/mocker/src/automocker.ts
+++ b/packages/mocker/src/automocker.ts
@@ -75,7 +75,7 @@ export function mockObject(
       const isFunction
         = type.includes('Function') && typeof value === 'function'
       if (
-        (!isFunction || value.__isMockFunction)
+        (!isFunction || value._isMockFunction)
         && type !== 'Object'
         && type !== 'Module'
       ) {
diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts
index 2de229bb9908..4f3180224621 100644
--- a/packages/spy/src/index.ts
+++ b/packages/spy/src/index.ts
@@ -413,6 +413,7 @@ export type Mocked<T> = {
       : T[P];
 } & T
 
+const vitestSpy = Symbol.for('vitest.spy')
 export const mocks: Set<MockInstance> = new Set()
 
 export function isMockFunction(fn: any): fn is MockInstance {
@@ -421,6 +422,20 @@ export function isMockFunction(fn: any): fn is MockInstance {
   )
 }
 
+function getSpy(
+  obj: unknown,
+  method: keyof any,
+  accessType?: 'get' | 'set',
+): MockInstance | undefined {
+  const desc = Object.getOwnPropertyDescriptor(obj, method)
+  if (desc) {
+    const fn = desc[accessType ?? 'value']
+    if (typeof fn === 'function' && vitestSpy in fn) {
+      return fn
+    }
+  }
+}
+
 export function spyOn<T, S extends Properties<Required<T>>>(
   obj: T,
   methodName: S,
@@ -450,6 +465,11 @@ export function spyOn<T, K extends keyof T>(
   } as const
   const objMethod = accessType ? { [dictionary[accessType]]: method } : method
 
+  const currentStub = getSpy(obj, method, accessType)
+  if (currentStub) {
+    return currentStub
+  }
+
   const stub = tinyspy.internalSpyOn(obj, objMethod as any)
 
   return enhanceSpy(stub) as MockInstance
@@ -523,6 +543,11 @@ function enhanceSpy<T extends Procedure>(
 
   let name: string = (stub as any).name
 
+  Object.defineProperty(stub, vitestSpy, {
+    value: true,
+    enumerable: false,
+  })
+
   stub.getMockName = () => name || 'vi.fn()'
   stub.mockName = (n) => {
     name = n
diff --git a/test/core/test/jest-mock.test.ts b/test/core/test/jest-mock.test.ts
index a6a888c6cfd1..32c699601c82 100644
--- a/test/core/test/jest-mock.test.ts
+++ b/test/core/test/jest-mock.test.ts
@@ -363,6 +363,31 @@ describe('jest mock compat layer', () => {
     expect(obj.property).toBe(true)
   })
 
+  it('spyOn returns the same spy twice', () => {
+    const obj = {
+      method() {
+        return 'original'
+      },
+    }
+
+    const spy1 = vi.spyOn(obj, 'method').mockImplementation(() => 'mocked')
+    const spy2 = vi.spyOn(obj, 'method')
+
+    expect(vi.isMockFunction(obj.method)).toBe(true)
+    expect(obj.method()).toBe('mocked')
+    expect(spy1).toBe(spy2)
+
+    spy2.mockImplementation(() => 'mocked2')
+
+    expect(obj.method()).toBe('mocked2')
+
+    spy2.mockRestore()
+
+    expect(obj.method()).toBe('original')
+    expect(vi.isMockFunction(obj.method)).toBe(false)
+    expect(obj.method).not.toBe(spy1)
+  })
+
   it('should spy on property setter (2), and mockReset should not restore original descriptor', () => {
     const obj = {
       _property: false,
diff --git a/test/core/test/mocked.test.ts b/test/core/test/mocked.test.ts
index 47bfe9e24106..175b204a3cd9 100644
--- a/test/core/test/mocked.test.ts
+++ b/test/core/test/mocked.test.ts
@@ -78,6 +78,9 @@ describe('mocked classes', () => {
 
     expect(MockedC.prototype.doSomething).toHaveBeenCalledOnce()
     expect(MockedC.prototype.doSomething).not.toHaveReturnedWith('A')
+
+    vi.mocked(instance.doSomething).mockRestore()
+    expect(instance.doSomething()).not.toBe('A')
   })
 
   test('should mock getters', () => {