fix(spy): copy static properties if spy is initialised with vi.fn(), fix types for vi.spyOn(obj, class) (#8956)

Author: sheremet-va

packages/spy/src/index.ts

@@ -194,6 +194,12 @@ export function createMockInstance(options: MockInstanceOption = {}): Mock<Proce
 export function fn<T extends Procedure | Constructable = Procedure>(
   originalImplementation?: T,
 ): Mock<T> {
+  // if the function is already a mock, just return the same function,
+  // simillarly to how vi.spyOn() works
+  if (originalImplementation != null && isMockFunction(originalImplementation)) {
+    return originalImplementation as Mock<T>
+  }
+
   return createMockInstance({
     // we pass this down so getMockImplementation() always returns the value
     mockImplementation: originalImplementation,
@@ -216,11 +222,9 @@ export function spyOn<T extends object, G extends Properties<Required<T>>>(
 export function spyOn<T extends object, M extends Classes<Required<T>> | Methods<Required<T>>>(
   object: T,
   key: M
-): Required<T>[M] extends { new (...args: infer A): infer R }
-  ? Mock<{ new (...args: A): R }>
-  : Required<T>[M] extends Procedure
-    ? Mock<Required<T>[M]>
-    : never
+): Required<T>[M] extends Constructable | Procedure
+  ? Mock<Required<T>[M]>
+  : never
 export function spyOn<T extends object, K extends keyof T>(
   object: T,
   key: K,
@@ -374,13 +378,15 @@ function createMock(
     prototypeState,
     prototypeConfig,
     keepMembersImplementation,
+    mockImplementation,
     prototypeMembers = [],
   }: MockInstanceOption & {
     state: MockContext
     config: MockConfig
   },
 ) {
-  const original = config.mockOriginal
+  const original = config.mockOriginal // init with vi.spyOn(obj, 'Klass')
+  const pseudoOriginal = mockImplementation // init with vi.fn(Klass)
   const name = (mockName || original?.name || 'Mock') as string
   const namedObject: Record<string, Mock<Procedure | Constructable>> = {
     // to keep the name of the function intact
@@ -499,8 +505,9 @@ function createMock(
     }) as Mock,
   }
   const mock = namedObject[name] as Mock<Procedure | Constructable>
-  if (original) {
-    copyOriginalStaticProperties(mock, original)
+  const copyPropertiesFrom = original || pseudoOriginal
+  if (copyPropertiesFrom) {
+    copyOriginalStaticProperties(mock, copyPropertiesFrom)
   }
   return mock
 }

test/core/test/mocking/vi-fn.test.ts

@@ -20,6 +20,37 @@ test('vi.fn().mock cannot be overriden', () => {
   }).toThrowError()
 })
 
+describe('vi.fn() copies static properties', () => {
+  test('vi.fn() copies static properties from functions', () => {
+    function Example() {}
+    Example.HELLO_WORLD = true
+
+    const spy = vi.fn(Example)
+    expect(Example.HELLO_WORLD).toBe(true)
+    expect(spy.HELLO_WORLD).toBe(true)
+  })
+
+  test('vi.fn() copies static properties from classes', () => {
+    class Example {
+      static HELLO_WORLD = true
+    }
+
+    const spy = vi.fn(Example)
+    expect(Example.HELLO_WORLD).toBe(true)
+    expect(spy.HELLO_WORLD).toBe(true)
+  })
+
+  test('vi.fn() ignores "node.js.promisify" symbol', () => {
+    const promisifySymbol = Symbol.for('nodejs.util.promisify.custom')
+    class Example {
+      static [promisifySymbol] = () => Promise.resolve(42)
+    }
+
+    const spy = vi.fn(Example)
+    expect(spy[promisifySymbol]).toBe(undefined)
+  })
+})
+
 describe('fn.length is consistent', () => {
   test('vi.fn() has correct length', () => {
     const fn0 = vi.fn(() => {})

test/core/test/mocking/vi-spyOn.test.ts

@@ -18,6 +18,47 @@ test('vi.fn() has correct length', () => {
   expect(fn3.length).toBe(3)
 })
 
+describe('vi.spyOn() copies static properties', () => {
+  test('vi.spyOn() 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.HELLO_WORLD).toBe(true)
+  })
+
+  test('vi.spyOn() 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.HELLO_WORLD).toBe(true)
+  })
+
+  test('vi.spyOn() ignores node.js.promisify symbol', () => {
+    const promisifySymbol = Symbol.for('nodejs.util.promisify.custom')
+    class Example {
+      static [promisifySymbol] = () => Promise.resolve(42)
+    }
+    const obj = { Example }
+
+    const spy = vi.spyOn(obj, 'Example')
+
+    expect(spy[promisifySymbol]).toBe(undefined)
+  })
+})
+
 describe('vi.spyOn() state', () => {
   test('vi.spyOn() spies on an object and tracks the calls', () => {
     const object = createObject()

test/core/test/spy.test.ts

@@ -30,37 +30,4 @@ 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('ignores node.js.promisify symbol', () => {
-    const promisifySymbol = Symbol.for('nodejs.util.promisify.custom')
-    class Example {
-      static [promisifySymbol] = () => Promise.resolve(42)
-    }
-    const obj = { Example }
-    const spy = vi.spyOn(obj, 'Example')
-    expect((spy as any)[promisifySymbol]).toBe(undefined)
-  })
 })