feat(types): add partial flag to mocked TypeScript helper (#1739) (#1742

)

Co-authored-by: jfrs <jfrs@jfrs.me>

docs/api/index.md

@@ -1830,9 +1830,11 @@ Vitest provides utility functions to help you out through it's **vi** helper. Yo
 ### vi.mocked
 
 - **Type**: `<T>(obj: T, deep?: boolean) => MaybeMockedDeep<T>`
+- **Type**: `<T>(obj: T, options?: { partial?: boolean; deep?: boolean }) => MaybePartiallyMockedDeep<T>`
 
   Type helper for TypeScript. In reality just returns the object that was passed.
 
+  When `partial` is `true` it will expect a `Partial<T>` as a return value.
   ```ts
   import example from './example'
   vi.mock('./example')

packages/vitest/src/integrations/spy.ts

@@ -61,6 +61,10 @@ export interface Mock<TArgs extends any[] = any, TReturns = any> extends SpyInst
   new (...args: TArgs): TReturns
   (...args: TArgs): TReturns
 }
+export interface PartialMock<TArgs extends any[] = any, TReturns = any> extends SpyInstance<TArgs, Partial<TReturns>> {
+  new (...args: TArgs): TReturns
+  (...args: TArgs): TReturns
+}
 
 export type MaybeMockedConstructor<T> = T extends new (
   ...args: Array<any>
@@ -70,7 +74,11 @@ export type MaybeMockedConstructor<T> = T extends new (
 export type MockedFunction<T extends Procedure> = Mock<Parameters<T>, ReturnType<T>> & {
   [K in keyof T]: T[K];
 }
+export type PartiallyMockedFunction<T extends Procedure> = PartialMock<Parameters<T>, ReturnType<T>> & {
+  [K in keyof T]: T[K];
+}
 export type MockedFunctionDeep<T extends Procedure> = Mock<Parameters<T>, ReturnType<T>> & MockedObjectDeep<T>
+export type PartiallyMockedFunctionDeep<T extends Procedure> = PartialMock<Parameters<T>, ReturnType<T>> & MockedObjectDeep<T>
 export type MockedObject<T> = MaybeMockedConstructor<T> & {
   [K in Methods<T>]: T[K] extends Procedure
     ? MockedFunction<T[K]>
@@ -88,12 +96,24 @@ export type MaybeMockedDeep<T> = T extends Procedure
     ? MockedObjectDeep<T>
     : T
 
+export type MaybePartiallyMockedDeep<T> = T extends Procedure
+  ? PartiallyMockedFunctionDeep<T>
+  : T extends object
+    ? MockedObjectDeep<T>
+    : T
+
 export type MaybeMocked<T> = T extends Procedure
   ? MockedFunction<T>
   : T extends object
     ? MockedObject<T>
     : T
 
+export type MaybePartiallyMocked<T> = T extends Procedure
+  ? PartiallyMockedFunction<T>
+  : T extends object
+    ? MockedObject<T>
+    : T
+
 interface Constructable {
   new (...args: any[]): any
 }

packages/vitest/src/integrations/vi.ts

@@ -3,7 +3,7 @@ import { parseStacktrace } from '../utils/source-map'
 import type { VitestMocker } from '../runtime/mocker'
 import { getWorkerState, resetModules, setTimeout } from '../utils'
 import { FakeTimers } from './mock/timers'
-import type { EnhancedSpy, MaybeMocked, MaybeMockedDeep } from './spy'
+import type { EnhancedSpy, MaybeMocked, MaybeMockedDeep, MaybePartiallyMocked, MaybePartiallyMockedDeep } from './spy'
 import { fn, isMockFunction, spies, spyOn } from './spy'
 
 class VitestUtils {
@@ -173,6 +173,8 @@ class VitestUtils {
 
   /**
    * Type helpers for TypeScript. In reality just returns the object that was passed.
+   *
+   * When `partial` is `true` it will expect a `Partial<T>` as a return value.
    * @example
    * import example from './example'
    * vi.mock('./example')
@@ -186,10 +188,15 @@ class VitestUtils {
    * })
    * @param item Anything that can be mocked
    * @param deep If the object is deeply mocked
+   * @param options If the object is partially or deeply mocked
    */
   public mocked<T>(item: T, deep?: false): MaybeMocked<T>
   public mocked<T>(item: T, deep: true): MaybeMockedDeep<T>
-  public mocked<T>(item: T, _deep = false): MaybeMocked<T> | MaybeMockedDeep<T> {
+  public mocked<T>(item: T, options: { partial?: false; deep?: false }): MaybeMocked<T>
+  public mocked<T>(item: T, options: { partial?: false; deep: true }): MaybeMockedDeep<T>
+  public mocked<T>(item: T, options: { partial: true; deep?: false }): MaybePartiallyMocked<T>
+  public mocked<T>(item: T, options: { partial: true; deep: true }): MaybePartiallyMockedDeep<T>
+  public mocked<T>(item: T, _options = {}): MaybeMocked<T> {
     return item as any
   }
 

test/core/test/vi.spec.ts

@@ -40,6 +40,30 @@ describe('testing vi utils', () => {
     expectType<MockedFunction<() => boolean>>(vi.fn())
   })
 
+  test('vi partial mocked', () => {
+    interface FooBar {
+      foo: () => void
+      bar: () => boolean
+      baz: string
+    }
+
+    type FooBarFactory = () => FooBar
+
+    const mockFactory: FooBarFactory = vi.fn()
+
+    vi.mocked(mockFactory, { partial: true }).mockReturnValue({
+      foo: vi.fn(),
+    })
+
+    vi.mocked(mockFactory, { partial: true, deep: false }).mockReturnValue({
+      bar: vi.fn(),
+    })
+
+    vi.mocked(mockFactory, { partial: true, deep: true }).mockReturnValue({
+      baz: 'baz',
+    })
+  })
+
   // TODO: it's unstable in CI, skip until resolved
   test.skip('loads unloaded module', async () => {
     let mod: any