fix(spy): copy over static properties from the function (#7780)

Author: sheremet-va

packages/spy/package.json

@@ -33,6 +33,6 @@
     "dev": "rollup -c --watch"
   },
   "dependencies": {
-    "tinyspy": "^3.0.2"
+    "tinyspy": "^4.0.3"
   }
 }

packages/spy/src/index.ts

@@ -145,6 +145,14 @@ export interface MockContext<T extends Procedure> {
    * @see https://vitest.dev/api/mock#mock-lastcall
    */
   lastCall: Parameters<T> | undefined
+  /** @internal */
+  _state: (state?: InternalState) => InternalState
+}
+
+interface InternalState {
+  implementation: Procedure | undefined
+  onceImplementations: Procedure[]
+  implementationChangedTemporarily: boolean
 }
 
 type Procedure = (...args: any[]) => any
@@ -412,7 +420,7 @@ export type Mocked<T> = {
       : T[P];
 } & T
 
-export const mocks: Set<MockInstance> = new Set()
+export const mocks: Set<MockInstance<any>> = new Set()
 
 export function isMockFunction(fn: any): fn is MockInstance {
   return (
@@ -449,23 +457,40 @@ export function spyOn<T, K extends keyof T>(
   } as const
   const objMethod = accessType ? { [dictionary[accessType]]: method } : method
 
+  let state: InternalState | undefined
+
+  const descriptor = getDescriptor(obj, method)
+  const fn = descriptor && descriptor[accessType || 'value']
+
+  // inherit implementations if it was already mocked
+  if (isMockFunction(fn)) {
+    state = fn.mock._state()
+  }
+
   const stub = tinyspy.internalSpyOn(obj, objMethod as any)
+  const spy = enhanceSpy(stub) as MockInstance
+
+  if (state) {
+    spy.mock._state(state)
+  }
 
-  return enhanceSpy(stub) as MockInstance
+  return spy
 }
 
 let callOrder = 0
 
 function enhanceSpy<T extends Procedure>(
   spy: SpyInternalImpl<Parameters<T>, ReturnType<T>>,
 ): MockInstance<T> {
-  type TArgs = Parameters<T>
   type TReturns = ReturnType<T>
 
   const stub = spy as unknown as MockInstance<T>
 
   let implementation: T | undefined
 
+  let onceImplementations: T[] = []
+  let implementationChangedTemporarily = false
+
   let instances: any[] = []
   let contexts: any[] = []
   let invocations: number[] = []
@@ -502,11 +527,20 @@ function enhanceSpy<T extends Procedure>(
     get lastCall() {
       return state.calls[state.calls.length - 1]
     },
+    _state(state) {
+      if (state) {
+        implementation = state.implementation as T
+        onceImplementations = state.onceImplementations as T[]
+        implementationChangedTemporarily = state.implementationChangedTemporarily
+      }
+      return {
+        implementation,
+        onceImplementations,
+        implementationChangedTemporarily,
+      }
+    },
   }
 
-  let onceImplementations: ((...args: TArgs) => TReturns)[] = []
-  let implementationChangedTemporarily = false
-
   function mockCall(this: unknown, ...args: any) {
     instances.push(this)
     contexts.push(this)
@@ -582,7 +616,7 @@ function enhanceSpy<T extends Procedure>(
 
     const result = cb()
 
-    if (result instanceof Promise) {
+    if (typeof result === 'object' && result && typeof result.then === 'function') {
       return result.then(() => {
         reset()
         return stub
@@ -639,3 +673,21 @@ export function fn<T extends Procedure = Procedure>(
 
   return enhancedSpy as any
 }
+
+function getDescriptor(
+  obj: any,
+  method: string | symbol | number,
+): PropertyDescriptor | undefined {
+  const objDescriptor = Object.getOwnPropertyDescriptor(obj, method)
+  if (objDescriptor) {
+    return objDescriptor
+  }
+  let currentProto = Object.getPrototypeOf(obj)
+  while (currentProto !== null) {
+    const descriptor = Object.getOwnPropertyDescriptor(currentProto, method)
+    if (descriptor) {
+      return descriptor
+    }
+    currentProto = Object.getPrototypeOf(currentProto)
+  }
+}

pnpm-lock.yaml

@@ -32,7 +32,7 @@
     "sweetalert2": "^11.22.0",
     "temporal-polyfill": "~0.3.0",
     "tinyrainbow": "catalog:",
-    "tinyspy": "^1.1.1",
+    "tinyspy": "^4.0.1",
     "url": "^0.11.4",
     "vite-node": "workspace:*",
     "vitest": "workspace:*",

test/core/package.json

@@ -363,6 +363,33 @@ describe('jest mock compat layer', () => {
     expect(obj.property).toBe(true)
   })
 
+  it('respyin on a spy resets the counter', () => {
+    const obj = {
+      method() {
+        return 'original'
+      },
+    }
+    vi.spyOn(obj, 'method')
+    obj.method()
+    expect(obj.method).toHaveBeenCalledTimes(1)
+    vi.spyOn(obj, 'method')
+    obj.method()
+    expect(obj.method).toHaveBeenCalledTimes(1)
+  })
+
+  it('spyOn on the getter multiple times', () => {
+    const obj = {
+      get getter() {
+        return 'original'
+      },
+    }
+
+    vi.spyOn(obj, 'getter', 'get').mockImplementation(() => 'mocked')
+    vi.spyOn(obj, 'getter', 'get')
+
+    expect(obj.getter).toBe('mocked')
+  })
+
   it('spyOn multiple times', () => {
     const obj = {
       method() {
@@ -383,9 +410,9 @@ describe('jest mock compat layer', () => {
 
     spy2.mockRestore()
 
-    expect(obj.method()).toBe('mocked')
-    expect(vi.isMockFunction(obj.method)).toBe(true)
-    expect(obj.method).toBe(spy1)
+    expect(obj.method()).toBe('original')
+    expect(vi.isMockFunction(obj.method)).toBe(false)
+    expect(obj.method).not.toBe(spy1)
 
     spy1.mockRestore()
     expect(vi.isMockFunction(obj.method)).toBe(false)
@@ -560,11 +587,104 @@ describe('jest mock compat layer', () => {
     expect(fn.getMockImplementation()).toBe(temporaryMockImplementation)
   })
 
+  it('keeps the descriptor the same as the original one when restoring', () => {
+    class Foo {
+      f() {
+        return 'original'
+      }
+    }
+
+    // initially there's no own properties
+    const foo = new Foo()
+    expect(foo.f()).toMatchInlineSnapshot(`"original"`)
+    expect(Object.getOwnPropertyDescriptors(foo)).toMatchInlineSnapshot(`{}`)
+
+    // mocked function in own property
+    const spy = vi.spyOn(foo, 'f').mockImplementation(() => 'mocked')
+    expect(foo.f()).toMatchInlineSnapshot(`"mocked"`)
+    expect(Object.getOwnPropertyDescriptors(foo)).toMatchInlineSnapshot(`
+    {
+      "f": {
+        "configurable": true,
+        "enumerable": false,
+        "value": [MockFunction f] {
+          "calls": [
+            [],
+          ],
+          "results": [
+            {
+              "type": "return",
+              "value": "mocked",
+            },
+          ],
+        },
+        "writable": true,
+      },
+    }
+  `)
+
+    // probably original prototype method is not moved to own property
+    spy.mockRestore()
+    expect(foo.f()).toMatchInlineSnapshot(`"original"`)
+    expect(Object.getOwnPropertyDescriptors(foo)).toMatchInlineSnapshot(`{}`)
+  })
+
+  it('mocks inherited methods', () => {
+    class Bar {
+      _bar = 'bar'
+      get bar(): string {
+        return this._bar
+      }
+
+      set bar(bar: string) {
+        this._bar = bar
+      }
+    }
+    class Foo extends Bar {}
+    const foo = new Foo()
+    vi.spyOn(foo, 'bar', 'get').mockImplementation(() => 'foo')
+    expect(foo.bar).toEqual('foo')
+    // foo.bar setter is inherited from Bar, so we can set it
+    expect(() => {
+      foo.bar = 'baz'
+    }).not.toThrowError()
+    expect(foo.bar).toEqual('foo')
+  })
+
+  it('mocks inherited overridden methods', () => {
+    class Bar {
+      _bar = 'bar'
+      get bar(): string {
+        return this._bar
+      }
+
+      set bar(bar: string) {
+        this._bar = bar
+      }
+    }
+    class Foo extends Bar {
+      get bar(): string {
+        return `${super.bar}-foo`
+      }
+    }
+    const foo = new Foo()
+    expect(foo.bar).toEqual('bar-foo')
+    vi.spyOn(foo, 'bar', 'get').mockImplementation(() => 'foo')
+    expect(foo.bar).toEqual('foo')
+    // foo.bar setter is not inherited from Bar
+    expect(() => {
+      // @ts-expect-error bar is readonly
+      foo.bar = 'baz'
+    }).toThrowError()
+    expect(foo.bar).toEqual('foo')
+  })
+
   describe('is disposable', () => {
     describe.runIf(Symbol.dispose)('in environments supporting it', () => {
       it('has dispose property', () => {
         expect(vi.fn()[Symbol.dispose]).toBeTypeOf('function')
       })
+
       it('calls mockRestore when disposing', () => {
         const fn = vi.fn()
         const restoreSpy = vi.spyOn(fn, 'mockRestore')
@@ -573,10 +693,12 @@ describe('jest mock compat layer', () => {
         }
         expect(restoreSpy).toHaveBeenCalled()
       })
+
       it('allows disposal when using mockImplementation', () => {
         expect(vi.fn().mockImplementation(() => {})[Symbol.dispose]).toBeTypeOf('function')
       })
     })
+
     describe.skipIf(Symbol.dispose)('in environments not supporting it', () => {
       it('does not have dispose property', () => {
         expect(vi.fn()[Symbol.dispose]).toBeUndefined()

test/core/test/jest-mock.test.ts

@@ -29,4 +29,27 @@ describe('spyOn', () => {
 
     expect(hw.hello()).toEqual('hello world')
   })
+
+  test('spying copies properties from functions', () => {
+    function a() {}
+    a.HELLO_WORLD = true
+    const obj = {
+      a,
+    }
+    const spy = vi.spyOn(obj, 'a')
+    expect(obj.a.HELLO_WORLD).toBe(true)
+    expect((spy as any).HELLO_WORLD).toBe(true)
+  })
+
+  test('spying copies properties from classes', () => {
+    class A {
+      static HELLO_WORLD = true
+    }
+    const obj = {
+      A,
+    }
+    const spy = vi.spyOn(obj, 'A')
+    expect(obj.A.HELLO_WORLD).toBe(true)
+    expect((spy as any).HELLO_WORLD).toBe(true)
+  })
 })

test/core/test/spy.test.ts

@@ -71,7 +71,7 @@ describe('import durations', () => {
 
     const throwsFile = resolve(root, 'import-durations-25ms-throws.ts')
 
-    expect(file.importDurations?.[throwsFile]?.totalTime).toBeGreaterThanOrEqual(25)
-    expect(file.importDurations?.[throwsFile]?.selfTime).toBeGreaterThanOrEqual(25)
+    expect(file.importDurations?.[throwsFile]?.totalTime).toBeGreaterThanOrEqual(24)
+    expect(file.importDurations?.[throwsFile]?.selfTime).toBeGreaterThanOrEqual(24)
   })
 })