feat(vitest): add return type and promisable mockFactory (#6139)

syi0808 · sheremet-va · web-flow · commit f5e0b9877f23 · 2024-08-12T16:14:12.000+02:00
Co-authored-by: Vladimir Sheremet &lt;sleuths.slews0s@icloud.com&gt;
diff --git a/docs/api/vi.md b/docs/api/vi.md
@@ -17,7 +17,7 @@ This section describes the API that you can use when [mocking a module](/guide/m
 ### vi.mock
 
 - **Type**: `(path: string, factory?: (importOriginal: () => unknown) => unknown) => void`
-- **Type**: `<T>(path: Promise<T>, factory?: (importOriginal: () => T) => unknown) => void`
+- **Type**: `<T>(path: Promise<T>, factory?: (importOriginal: () => T) => T | Promise<T>) => void`
 
 Substitutes all imported modules from provided `path` with another module. You can use configured Vite aliases inside a path. The call to `vi.mock` is hoisted, so it doesn't matter where you call it. It will always be executed before all imports. If you need to reference some variables outside of its scope, you can define them inside [`vi.hoisted`](#vi-hoisted) and reference them inside `vi.mock`.
 
@@ -31,7 +31,9 @@ Vitest will not mock modules that were imported inside a [setup file](/config/#s
 
 If `factory` is defined, all imports will return its result. Vitest calls factory only once and caches results for all subsequent imports until [`vi.unmock`](#vi-unmock) or [`vi.doUnmock`](#vi-dounmock) is called.
 
-Unlike in `jest`, the factory can be asynchronous. You can use [`vi.importActual`](#vi-importactual) or a helper with the factory passed in as the first argument, and get the original module inside. Vitest also supports a module promise instead of a string in `vi.mock` method for better IDE support (when file is moved, path will be updated, `importOriginal` also inherits the type automatically).
+Unlike in `jest`, the factory can be asynchronous. You can use [`vi.importActual`](#vi-importactual) or a helper with the factory passed in as the first argument, and get the original module inside.
+
+Vitest also supports a module promise instead of a string in the `vi.mock` and `vi.doMock` methods for better IDE support. When the file is moved, the path will be updated, and `importOriginal` also inherits the type automatically. Using this signature will also enforce factory return type to be compatible with the original module (but every export is optional).
 
 ```ts twoslash
 // @filename: ./path/to/module.js
@@ -143,7 +145,7 @@ If there is no `__mocks__` folder or a factory provided, Vitest will import the
 ### vi.doMock
 
 - **Type**: `(path: string, factory?: (importOriginal: () => unknown) => unknown) => void`
-- **Type**: `<T>(path: Promise<T>, factory?: (importOriginal: () => T) => unknown) => void`
+- **Type**: `<T>(path: Promise<T>, factory?: (importOriginal: () => T) => T | Promise<T>) => void`
 
 The same as [`vi.mock`](#vi-mock), but it's not hoisted to the top of the file, so you can reference variables in the global file scope. The next [dynamic import](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/import) of the module will be mocked.
 
diff --git a/packages/vitest/src/types/mocker.ts b/packages/vitest/src/types/mocker.ts
@@ -1,6 +1,8 @@
+type Promisable<T> = T | Promise<T>
+
 export type MockFactoryWithHelper<M = unknown> = (
-  importOriginal: <T extends M>() => Promise<T>
-) => any
+  importOriginal: <T extends M = M>() => Promise<T>
+) => Promisable<Partial<M>>
 export type MockFactory = () => any
 
 export type MockMap = Map<string, Record<string, string | null | MockFactory>>
diff --git a/test/config/test/workers-option.test.ts b/test/config/test/workers-option.test.ts
@@ -4,6 +4,7 @@ import { getWorkersCountByPercentage } from 'vitest/src/utils/workers.js'
 import * as testUtils from '../../test-utils'
 
 vi.mock(import('node:os'), async importOriginal => ({
+  ...(await importOriginal()),
   default: {
     ...(await importOriginal()).default,
     availableParallelism: () => 10,
diff --git a/test/core/test/mocking/factory.test.ts b/test/core/test/mocking/factory.test.ts
@@ -54,7 +54,7 @@ vi.mock('../../src/mocks/log.ts', async () => {
     },
   }
 })
-
+// @ts-expect-error null is not allowed to mocked implementation
 vi.mock('../../src/mocks/default.ts', () => null)
 
 describe('mocking with factory', () => {

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ vi.mock('../../src/mocks/log.ts', async () => {`
`54`	`54`	`},`
`55`	`55`	`}`
`56`	`56`	`})`
`57`		`-`
	`57`	`+// @ts-expect-error null is not allowed to mocked implementation`
`58`	`58`	`vi.mock('../../src/mocks/default.ts', () => null)`
`59`	`59`
`60`	`60`	`describe('mocking with factory', () => {`