From 5ce2bb7b43d503fa0ed0e8da7f0f389ec7be4387 Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Mon, 9 Sep 2024 09:02:40 +0200
Subject: [PATCH] fix: reuse spy when calling spyOn (#47)

* fix: reuse spy when calling spyOn

* test: remove only
---
 src/internal.ts    |  5 +++++
 src/spyOn.ts       | 36 ++++++++++++++++++++++--------------
 test/index.test.ts | 20 ++++++++++++++++++++
 3 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/src/internal.ts b/src/internal.ts
index 317e3be..8db9385 100644
--- a/src/internal.ts
+++ b/src/internal.ts
@@ -138,6 +138,11 @@ export function createInternalSpy<A extends any[], R>(
 export function populateSpy<A extends any[], R>(spy: SpyInternal<A, R>) {
   const I = getInternalState(spy)
 
+  // already populated
+  if ('returns' in spy) {
+    return
+  }
+
   define(spy, 'returns', {
     get: () => I.results.map(([, r]) => r),
   })
diff --git a/src/spyOn.ts b/src/spyOn.ts
index 3c2682e..2c8b34b 100644
--- a/src/spyOn.ts
+++ b/src/spyOn.ts
@@ -3,6 +3,7 @@ import {
   populateSpy,
   spies,
   SpyImpl,
+  SpyInternal,
   SpyInternalImpl,
 } from './internal'
 import { assert, define, defineValue, isType } from './utils'
@@ -47,7 +48,7 @@ export function internalSpyOn<T, K extends string & keyof T>(
 
   let [accessName, accessType] = ((): [
     string | symbol | number,
-    'value' | 'get' | 'set'
+    'value' | 'get' | 'set',
   ] => {
     if (!isType('object', methodName)) {
       return [methodName, 'value']
@@ -97,12 +98,6 @@ export function internalSpyOn<T, K extends string & keyof T>(
     origin = obj[accessName as keyof T] as unknown as Procedure
   }
 
-  if (!mock) mock = origin
-
-  let fn = createInternalSpy(mock)
-  if (accessType === 'value') {
-    prototype(fn, origin)
-  }
   let reassign = (cb: any) => {
     let { value, ...desc } = originalDescriptor || {
       configurable: true,
@@ -118,13 +113,26 @@ export function internalSpyOn<T, K extends string & keyof T>(
     originalDescriptor
       ? define(obj, accessName, originalDescriptor)
       : reassign(origin)
-  const state = fn[S]
-  defineValue(state, 'restore', restore)
-  defineValue(state, 'getOriginal', () => (ssr ? origin() : origin))
-  defineValue(state, 'willCall', (newCb: Procedure) => {
-    state.impl = newCb
-    return fn
-  })
+
+  if (!mock) mock = origin
+
+  let fn: SpyInternal
+  if (origin && S in origin) {
+    fn = origin as SpyInternal
+  } else {
+    fn = createInternalSpy(mock)
+    if (accessType === 'value') {
+      prototype(fn, origin)
+    }
+
+    const state = fn[S]
+    defineValue(state, 'restore', restore)
+    defineValue(state, 'getOriginal', () => (ssr ? origin() : origin))
+    defineValue(state, 'willCall', (newCb: Procedure) => {
+      state.impl = newCb
+      return fn
+    })
+  }
 
   reassign(
     ssr
diff --git a/test/index.test.ts b/test/index.test.ts
index 18b29b0..02d0e8c 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -755,3 +755,23 @@ test('next in a row', () => {
   expect(cb()).toBe(3)
   expect(cb()).toBe(undefined)
 })
+
+test('spying twice and unspying restores original method', () => {
+  const obj = {
+    method: () => 1,
+  }
+  const spy1 = spyOn(obj, 'method').willCall(() => 2)
+  expect(obj.method()).toBe(2)
+
+  const spy2 = spyOn(obj, 'method')
+
+  expect(spy1).toBe(spy2)
+  expect(obj.method()).toBe(2)
+
+  spy2.willCall(() => 3)
+
+  expect(obj.method()).toBe(3)
+
+  spy2.restore()
+  expect(obj.method).not.toBe(spy1)
+})