From c18f0af54207785b3dd783a1ef2b0a1f5af5694b Mon Sep 17 00:00:00 2001
From: jfrs <4005764+jfrs@users.noreply.github.com>
Date: Tue, 2 Aug 2022 08:11:55 +0100
Subject: [PATCH] feat(types): add partial flag to mocked TypeScript helper
 (#1739) (#1742)

Co-authored-by: jfrs <jfrs@jfrs.me>
---
 docs/api/index.md                       |  2 ++
 packages/vitest/src/integrations/spy.ts | 20 ++++++++++++++++++++
 packages/vitest/src/integrations/vi.ts  | 11 +++++++++--
 test/core/test/vi.spec.ts               | 24 ++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/docs/api/index.md b/docs/api/index.md
index 7a695f64e707..d3c84c851e7c 100644
--- a/docs/api/index.md
+++ b/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')
diff --git a/packages/vitest/src/integrations/spy.ts b/packages/vitest/src/integrations/spy.ts
index 4231853d360c..b83d2c9a0f63 100644
--- a/packages/vitest/src/integrations/spy.ts
+++ b/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
 }
diff --git a/packages/vitest/src/integrations/vi.ts b/packages/vitest/src/integrations/vi.ts
index a459825109a8..086926d24bac 100644
--- a/packages/vitest/src/integrations/vi.ts
+++ b/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
   }
 
diff --git a/test/core/test/vi.spec.ts b/test/core/test/vi.spec.ts
index 63e03e540aa1..cd81d4f1fe89 100644
--- a/test/core/test/vi.spec.ts
+++ b/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