fix(spy): allow classes in vi.mocked utility (#8839)

Author: sheremet-va

packages/spy/src/types.ts

@@ -370,28 +370,36 @@ type PartialMaybePromise<T> = T extends Promise<Awaited<T>>
   ? Promise<Partial<Awaited<T>>>
   : Partial<T>
 
-export interface PartialMock<T extends Procedure = Procedure>
-  extends MockInstance<
-    (...args: Parameters<T>) => PartialMaybePromise<ReturnType<T>>
-  > {
-  new (...args: Parameters<T>): ReturnType<T>
-  (...args: Parameters<T>): ReturnType<T>
-}
+type PartialResultFunction<T> = T extends Constructable
+  ? ({
+    new (...args: ConstructorParameters<T>): InstanceType<T>
+  })
+  | ({
+    (this: InstanceType<T>, ...args: ConstructorParameters<T>): void
+  })
+  : T extends Procedure
+    ? (...args: Parameters<T>) => PartialMaybePromise<ReturnType<T>>
+    : T
+
+export interface PartialMock<T extends Procedure | Constructable = Procedure>
+  extends Mock<
+    PartialResultFunction<T extends Mock
+      ? NonNullable<ReturnType<T['getMockImplementation']>>
+      : T>
+  > {}
 
-export type MaybeMockedConstructor<T> = T extends new (
-  ...args: Array<any>
-) => infer R
-  ? Mock<(...args: ConstructorParameters<T>) => R>
+export type MaybeMockedConstructor<T> = T extends Constructable
+  ? Mock<T>
   : T
-export type MockedFunction<T extends Procedure> = Mock<T> & {
+export type MockedFunction<T extends Procedure | Constructable> = Mock<T> & {
   [K in keyof T]: T[K];
 }
-export type PartiallyMockedFunction<T extends Procedure> = PartialMock<T> & {
+export type PartiallyMockedFunction<T extends Procedure | Constructable> = PartialMock<T> & {
   [K in keyof T]: T[K];
 }
-export type MockedFunctionDeep<T extends Procedure> = Mock<T>
+export type MockedFunctionDeep<T extends Procedure | Constructable> = Mock<T>
   & MockedObjectDeep<T>
-export type PartiallyMockedFunctionDeep<T extends Procedure> = PartialMock<T>
+export type PartiallyMockedFunctionDeep<T extends Procedure | Constructable> = PartialMock<T>
   & MockedObjectDeep<T>
 export type MockedObject<T> = MaybeMockedConstructor<T> & {
   [K in Methods<T>]: T[K] extends Procedure ? MockedFunction<T[K]> : T[K];
@@ -400,25 +408,25 @@ export type MockedObjectDeep<T> = MaybeMockedConstructor<T> & {
   [K in Methods<T>]: T[K] extends Procedure ? MockedFunctionDeep<T[K]> : T[K];
 } & { [K in Properties<T>]: MaybeMockedDeep<T[K]> }
 
-export type MaybeMockedDeep<T> = T extends Procedure
+export type MaybeMockedDeep<T> = T extends Procedure | Constructable
   ? MockedFunctionDeep<T>
   : T extends object
     ? MockedObjectDeep<T>
     : T
 
-export type MaybePartiallyMockedDeep<T> = T extends Procedure
+export type MaybePartiallyMockedDeep<T> = T extends Procedure | Constructable
   ? PartiallyMockedFunctionDeep<T>
   : T extends object
     ? MockedObjectDeep<T>
     : T
 
-export type MaybeMocked<T> = T extends Procedure
+export type MaybeMocked<T> = T extends Procedure | Constructable
   ? MockedFunction<T>
   : T extends object
     ? MockedObject<T>
     : T
 
-export type MaybePartiallyMocked<T> = T extends Procedure
+export type MaybePartiallyMocked<T> = T extends Procedure | Constructable
   ? PartiallyMockedFunction<T>
   : T extends object
     ? MockedObject<T>
@@ -428,9 +436,7 @@ export interface Constructable {
   new (...args: any[]): any
 }
 
-export type MockedClass<T extends Constructable> = MockInstance<
-  (...args: ConstructorParameters<T>) => InstanceType<T>
-> & {
+export type MockedClass<T extends Constructable> = MockInstance<T> & {
   prototype: T extends { prototype: any } ? Mocked<T['prototype']> : never
 } & T
 

test/core/test/vi.spec.ts

@@ -2,7 +2,7 @@
  * @vitest-environment jsdom
  */
 
-import type { Mock, MockedFunction, MockedObject, MockInstance } from 'vitest'
+import type { Mock, Mocked, MockedFunction, MockedObject, MockInstance } from 'vitest'
 import { describe, expect, expectTypeOf, test, vi } from 'vitest'
 import { getWorkerState } from '../../../packages/vitest/src/runtime/utils'
 
@@ -56,33 +56,49 @@ describe('testing vi utils', () => {
       baz: string
     }
 
-    const mockFactory = vi.fn<() => FooBar>()
+    const mockFnFactory = vi.fn<() => FooBar>()
 
-    vi.mocked(mockFactory, { partial: true }).mockReturnValue({
+    vi.mocked(mockFnFactory, { partial: true }).mockReturnValue({
       foo: vi.fn(),
     })
 
-    vi.mocked(mockFactory, { partial: true, deep: false }).mockReturnValue({
+    vi.mocked(mockFnFactory, { partial: true, deep: false }).mockReturnValue({
       bar: vi.fn<FooBar['bar']>(),
     })
 
-    vi.mocked(mockFactory, { partial: true, deep: true }).mockReturnValue({
+    vi.mocked(mockFnFactory, { partial: true, deep: true }).mockReturnValue({
       baz: 'baz',
     })
 
-    const mockFactoryAsync = vi.fn<() => Promise<FooBar>>()
+    if (0) {
+      const mockFactory = (): FooBar => ({} as FooBar)
 
-    vi.mocked(mockFactoryAsync, { partial: true }).mockResolvedValue({
-      foo: vi.fn(),
-    })
+      vi.mocked(mockFactory, { partial: true }).mockReturnValue({
+        foo: vi.fn(),
+      })
 
-    vi.mocked(mockFactoryAsync, { partial: true, deep: false }).mockResolvedValue({
-      bar: vi.fn<FooBar['bar']>(),
-    })
+      vi.mocked(mockFactory, { partial: true, deep: false }).mockReturnValue({
+        bar: vi.fn<FooBar['bar']>(),
+      })
 
-    vi.mocked(mockFactoryAsync, { partial: true, deep: true }).mockResolvedValue({
-      baz: 'baz',
-    })
+      vi.mocked(mockFactory, { partial: true, deep: true }).mockReturnValue({
+        baz: 'baz',
+      })
+
+      const mockFactoryAsync = async (): Promise<FooBar> => ({} as FooBar)
+
+      vi.mocked(mockFactoryAsync, { partial: true }).mockResolvedValue({
+        foo: vi.fn(),
+      })
+
+      vi.mocked(mockFactoryAsync, { partial: true, deep: false }).mockResolvedValue({
+        bar: vi.fn<FooBar['bar']>(),
+      })
+
+      vi.mocked(mockFactoryAsync, { partial: true, deep: true }).mockResolvedValue({
+        baz: 'baz',
+      })
+    }
 
     function fetchSomething(): Promise<Response> {
       return fetch('https://vitest.dev/')
@@ -94,6 +110,33 @@ describe('testing vi utils', () => {
     }
   })
 
+  test('vi.mocked with classes', () => {
+    class Foo {
+      constructor(public readonly bar: string) {}
+
+      public getBar(): string {
+        return this.bar
+      }
+    }
+    class FooMock implements Mocked<Foo> {
+      readonly barMock: Mock<() => string> = vi.fn()
+
+      public get bar(): string {
+        return this.barMock()
+      }
+
+      public getBar: Mock<() => string> = vi
+        .fn()
+        .mockImplementation(() => this.barMock())
+    }
+
+    // type check only
+    if (0) {
+      vi.mocked(Foo).mockImplementation(FooMock)
+      vi.mocked(Foo).mockImplementation(Foo)
+    }
+  })
+
   test('vi.fn and Mock type', () => {
     // use case from https://github.com/vitest-dev/vitest/issues/4723#issuecomment-1851034249