Skip to content

Commit a99a14c

Browse files
sheremet-vasheremet-va
authored
feat: allow import statement as vi.mock path for better IDE support (#5690)
1 parent 8c96607 commit a99a14c

7 files changed

Lines changed: 217 additions & 14 deletions

File tree

‎docs/api/vi.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ This section describes the API that you can use when [mocking a module](/guide/m
1717
### vi.mock
1818

1919
- **Type**: `(path: string, factory?: (importOriginal: () => unknown) => unknown) => void`
20+
- **Type**: `<T>(path: Promise<T>, factory?: (importOriginal: () => T) => unknown) => void` <Version>2.0.0+</Version>
2021

2122
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`.
2223

@@ -64,6 +65,21 @@ vi.mock('./path/to/module.js', async (importOriginal) => {
6465
})
6566
```
6667

68+
Since 2.0.0, Vitest 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).
69+
70+
```ts
71+
vi.mock(import('./path/to/module.js'), async (importOriginal) => {
72+
const mod = await importOriginal() // type is inferred
73+
return {
74+
...mod,
75+
// replace some exports
76+
namedExport: vi.fn(),
77+
}
78+
})
79+
```
80+
81+
Under the hood, Vitest still operates on a string and not a module object.
82+
6783
::: warning
6884
`vi.mock` is hoisted (in other words, _moved_) to **top of the file**. It means that whenever you write it (be it inside `beforeEach` or `test`), it will actually be called before that.
6985

‎packages/vitest/src/integrations/vi.ts‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -184,15 +184,21 @@ export interface VitestUtils {
184184
* @param path Path to the module. Can be aliased, if your Vitest config supports it
185185
* @param factory Mocked module factory. The result of this function will be an exports object
186186
*/
187-
mock: (path: string, factory?: MockFactoryWithHelper) => void
187+
// eslint-disable-next-line ts/method-signature-style
188+
mock(path: string, factory?: MockFactoryWithHelper): void
189+
// eslint-disable-next-line ts/method-signature-style
190+
mock<T>(module: Promise<T>, factory?: MockFactoryWithHelper<T>): void
188191

189192
/**
190193
* Removes module from mocked registry. All calls to import will return the original module even if it was mocked before.
191194
*
192195
* This call is hoisted to the top of the file, so it will only unmock modules that were defined in `setupFiles`, for example.
193196
* @param path Path to the module. Can be aliased, if your Vitest config supports it
194197
*/
195-
unmock: (path: string) => void
198+
// eslint-disable-next-line ts/method-signature-style
199+
unmock(path: string): void
200+
// eslint-disable-next-line ts/method-signature-style
201+
unmock(module: Promise<unknown>): void
196202

197203
/**
198204
* Mocks every subsequent [dynamic import](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/import) call.
@@ -203,14 +209,20 @@ export interface VitestUtils {
203209
* @param path Path to the module. Can be aliased, if your Vitest config supports it
204210
* @param factory Mocked module factory. The result of this function will be an exports object
205211
*/
206-
doMock: (path: string, factory?: MockFactoryWithHelper) => void
212+
// eslint-disable-next-line ts/method-signature-style
213+
doMock(path: string, factory?: MockFactoryWithHelper): void
214+
// eslint-disable-next-line ts/method-signature-style
215+
doMock<T>(module: Promise<T>, factory?: MockFactoryWithHelper<T>): void
207216
/**
208217
* Removes module from mocked registry. All subsequent calls to import will return original module.
209218
*
210219
* Unlike [`vi.unmock`](https://vitest.dev/api/vi#vi-unmock), this method is not hoisted to the top of the file.
211220
* @param path Path to the module. Can be aliased, if your Vitest config supports it
212221
*/
213-
doUnmock: (path: string) => void
222+
// eslint-disable-next-line ts/method-signature-style
223+
doUnmock(path: string): void
224+
// eslint-disable-next-line ts/method-signature-style
225+
doUnmock(module: Promise<unknown>): void
214226

215227
/**
216228
* Imports module, bypassing all checks if it should be mocked.
@@ -476,7 +488,9 @@ function createVitest(): VitestUtils {
476488
return factory()
477489
},
478490

479-
mock(path: string, factory?: MockFactoryWithHelper) {
491+
mock(path: string | Promise<unknown>, factory?: MockFactoryWithHelper) {
492+
if (typeof path !== 'string')
493+
throw new Error(`vi.mock() expects a string path, but received a ${typeof path}`)
480494
const importer = getImporter()
481495
_mocker.queueMock(
482496
path,
@@ -486,11 +500,15 @@ function createVitest(): VitestUtils {
486500
)
487501
},
488502

489-
unmock(path: string) {
503+
unmock(path: string | Promise<unknown>) {
504+
if (typeof path !== 'string')
505+
throw new Error(`vi.unmock() expects a string path, but received a ${typeof path}`)
490506
_mocker.queueUnmock(path, getImporter())
491507
},
492508

493-
doMock(path: string, factory?: MockFactoryWithHelper) {
509+
doMock(path: string | Promise<unknown>, factory?: MockFactoryWithHelper) {
510+
if (typeof path !== 'string')
511+
throw new Error(`vi.doMock() expects a string path, but received a ${typeof path}`)
494512
const importer = getImporter()
495513
_mocker.queueMock(
496514
path,
@@ -500,7 +518,9 @@ function createVitest(): VitestUtils {
500518
)
501519
},
502520

503-
doUnmock(path: string) {
521+
doUnmock(path: string | Promise<unknown>) {
522+
if (typeof path !== 'string')
523+
throw new Error(`vi.doUnmock() expects a string path, but received a ${typeof path}`)
504524
_mocker.queueUnmock(path, getImporter())
505525
},
506526

‎packages/vitest/src/node/hoistMocks.ts‎

Lines changed: 61 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import MagicString from 'magic-string'
2-
import type { AwaitExpression, CallExpression, ExportDefaultDeclaration, ExportNamedDeclaration, Identifier, ImportDeclaration, VariableDeclaration, Node as _Node } from 'estree'
2+
import type { AwaitExpression, CallExpression, ExportDefaultDeclaration, ExportNamedDeclaration, Expression, Identifier, ImportDeclaration, ImportExpression, VariableDeclaration, Node as _Node } from 'estree'
33
import { findNodeAround } from 'acorn-walk'
4-
import type { PluginContext } from 'rollup'
4+
import type { PluginContext, ProgramNode } from 'rollup'
55
import { esmWalker } from '@vitest/utils/ast'
66
import type { Colors } from '@vitest/utils'
77
import { highlightCode } from '../utils/colors'
@@ -58,7 +58,7 @@ export function getBetterEnd(code: string, node: Node) {
5858
return end
5959
}
6060

61-
const regexpHoistable = /\b(vi|vitest)\s*\.\s*(mock|unmock|hoisted)\(/
61+
const regexpHoistable = /\b(vi|vitest)\s*\.\s*(mock|unmock|hoisted|doMock|doUnmock)\(/
6262
const hashbangRE = /^#!.*\n/
6363

6464
export function hoistMocks(code: string, id: string, parse: PluginContext['parse'], colors?: Colors) {
@@ -69,7 +69,7 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
6969

7070
const s = new MagicString(code)
7171

72-
let ast: any
72+
let ast: ProgramNode
7373
try {
7474
ast = parse(code)
7575
}
@@ -225,6 +225,24 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
225225
hoistedNodes.push(node)
226226
}
227227

228+
// vi.doMock(import('./path')) -> vi.doMock('./path')
229+
// vi.doMock(await import('./path')) -> vi.doMock('./path')
230+
if (methodName === 'doMock' || methodName === 'doUnmock') {
231+
const moduleInfo = node.arguments[0] as Positioned<Expression>
232+
let source: Positioned<Expression> | null = null
233+
if (moduleInfo.type === 'ImportExpression')
234+
source = moduleInfo.source as Positioned<Expression>
235+
if (moduleInfo.type === 'AwaitExpression' && moduleInfo.argument.type === 'ImportExpression')
236+
source = moduleInfo.argument.source as Positioned<Expression>
237+
if (source) {
238+
s.overwrite(
239+
moduleInfo.start,
240+
moduleInfo.end,
241+
s.slice(source.start, source.end),
242+
)
243+
}
244+
}
245+
228246
if (methodName === 'hoisted') {
229247
assertNotDefaultExport(node, 'Cannot export hoisted variable. You can control hoisting behavior by placing the import from this file first.')
230248

@@ -277,6 +295,14 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
277295
)
278296
}
279297

298+
function rewriteMockDynamicImport(nodeCode: string, moduleInfo: Positioned<ImportExpression>, expressionStart: number, expressionEnd: number, mockStart: number) {
299+
const source = moduleInfo.source as Positioned<Expression>
300+
const importPath = s.slice(source.start, source.end)
301+
const nodeCodeStart = expressionStart - mockStart
302+
const nodeCodeEnd = expressionEnd - mockStart
303+
return nodeCode.slice(0, nodeCodeStart) + importPath + nodeCode.slice(nodeCodeEnd)
304+
}
305+
280306
// validate hoistedNodes doesn't have nodes inside other nodes
281307
for (let i = 0; i < hoistedNodes.length; i++) {
282308
const node = hoistedNodes[i]
@@ -300,7 +326,37 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
300326
* import user from './user'
301327
* vi.mock('./mock.js', () => ({ getSession: vi.fn().mockImplementation(() => ({ user })) }))
302328
*/
303-
const nodeCode = s.slice(node.start, end)
329+
let nodeCode = s.slice(node.start, end)
330+
331+
// rewrite vi.mock(import('..')) into vi.mock('..')
332+
if (
333+
node.type === 'CallExpression'
334+
&& node.callee.type === 'MemberExpression'
335+
&& ((node.callee.property as Identifier).name === 'mock' || (node.callee.property as Identifier).name === 'unmock')
336+
) {
337+
const moduleInfo = node.arguments[0] as Positioned<Expression>
338+
// vi.mock(import('./path')) -> vi.mock('./path')
339+
if (moduleInfo.type === 'ImportExpression') {
340+
nodeCode = rewriteMockDynamicImport(
341+
nodeCode,
342+
moduleInfo,
343+
moduleInfo.start,
344+
moduleInfo.end,
345+
node.start,
346+
)
347+
}
348+
// vi.mock(await import('./path')) -> vi.mock('./path')
349+
if (moduleInfo.type === 'AwaitExpression' && moduleInfo.argument.type === 'ImportExpression') {
350+
nodeCode = rewriteMockDynamicImport(
351+
nodeCode,
352+
moduleInfo.argument as Positioned<ImportExpression>,
353+
moduleInfo.start,
354+
moduleInfo.end,
355+
node.start,
356+
)
357+
}
358+
}
359+
304360
s.remove(node.start, end)
305361
return `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
306362
}).join('')

‎packages/vitest/src/types/mocker.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
export type MockFactoryWithHelper = (importOriginal: <T = unknown>() => Promise<T>) => any
1+
export type MockFactoryWithHelper<M = unknown> = (importOriginal: <T extends M>() => Promise<T>) => any
22
export type MockFactory = () => any
33

44
export type MockMap = Map<string, Record<string, string | null | MockFactory>>

‎test/core/test/injector-mock.test.ts‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1201,6 +1201,104 @@ await vi
12011201
1234;"
12021202
`)
12031203
})
1204+
1205+
test('handles dynamic import as the first argument', () => {
1206+
expect(
1207+
hoistSimpleCode(`
1208+
vi.mock(import('./path'))
1209+
vi.mock(import(somePath))
1210+
vi.mock(import(\`./path\`))
1211+
1212+
vi.mock(import('./path'));
1213+
vi.mock(import(somePath));
1214+
vi.mock(import(\`./path\`));
1215+
1216+
vi.mock(await import('./path'))
1217+
vi.mock(await import(somePath))
1218+
vi.mock(await import(\`./path\`))
1219+
1220+
vi.mock(await import('./path'));
1221+
vi.mock(await import(somePath));
1222+
vi.mock(await import(\`./path\`));
1223+
1224+
vi.mock(import('./path'), () => {})
1225+
vi.mock(import(somePath), () => {})
1226+
vi.mock(import(\`./path\`), () => {})
1227+
1228+
vi.mock(await import('./path'), () => {})
1229+
vi.mock(await import(somePath), () => {})
1230+
vi.mock(await import(\`./path\`), () => {})
1231+
1232+
vi.mock(import('./path'), () => {});
1233+
vi.mock(import(somePath), () => {});
1234+
vi.mock(import(\`./path\`), () => {});
1235+
1236+
vi.mock(await import('./path'), () => {});
1237+
vi.mock(await import(somePath), () => {});
1238+
vi.mock(await import(\`./path\`), () => {});
1239+
`),
1240+
).toMatchInlineSnapshot(`
1241+
"if (typeof globalThis.vi === "undefined" && typeof globalThis.vitest === "undefined") { throw new Error("There are some problems in resolving the mocks API.\\nYou may encounter this issue when importing the mocks API from another module other than 'vitest'.\\nTo fix this issue you can either:\\n- import the mocks API directly from 'vitest'\\n- enable the 'globals' options") }
1242+
vi.mock('./path')
1243+
vi.mock(somePath)
1244+
vi.mock(\`./path\`)
1245+
vi.mock('./path');
1246+
vi.mock(somePath);
1247+
vi.mock(\`./path\`);
1248+
vi.mock('./path')
1249+
vi.mock(somePath)
1250+
vi.mock(\`./path\`)
1251+
vi.mock('./path');
1252+
vi.mock(somePath);
1253+
vi.mock(\`./path\`);
1254+
vi.mock('./path', () => {})
1255+
vi.mock(somePath, () => {})
1256+
vi.mock(\`./path\`, () => {})
1257+
vi.mock('./path', () => {})
1258+
vi.mock(somePath, () => {})
1259+
vi.mock(\`./path\`, () => {})
1260+
vi.mock('./path', () => {});
1261+
vi.mock(somePath, () => {});
1262+
vi.mock(\`./path\`, () => {});
1263+
vi.mock('./path', () => {});
1264+
vi.mock(somePath, () => {});
1265+
vi.mock(\`./path\`, () => {});"
1266+
`)
1267+
})
1268+
1269+
test.only('handles import in vi.do* methods', () => {
1270+
expect(
1271+
hoistSimpleCode(`
1272+
vi.doMock(import('./path'))
1273+
vi.doMock(import(\`./path\`))
1274+
vi.doMock(import('./path'));
1275+
1276+
beforeEach(() => {
1277+
vi.doUnmock(import('./path'))
1278+
vi.doMock(import('./path'))
1279+
})
1280+
1281+
test('test', async () => {
1282+
vi.doMock(import(dynamicName))
1283+
await import(dynamicName)
1284+
})
1285+
`),
1286+
).toMatchInlineSnapshot(`
1287+
"vi.doMock('./path')
1288+
vi.doMock(\`./path\`)
1289+
vi.doMock('./path');
1290+
1291+
beforeEach(() => {
1292+
vi.doUnmock('./path')
1293+
vi.doMock('./path')
1294+
})
1295+
1296+
test('test', async () => {
1297+
vi.doMock(dynamicName)
1298+
await import(dynamicName)
1299+
})"
1300+
`)
1301+
})
12041302
})
12051303

12061304
describe('throws an error when nodes are incompatible', () => {

‎test/core/test/mocked-import-circular.test.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
import { expect, it, vi } from 'vitest'
2+
3+
// The order of the two imports here matters: B before A
4+
import { circularB } from '../src/circularB'
5+
import { circularA } from '../src/circularA'
6+
7+
vi.mock(import('../src/circularB'))
8+
9+
it('circular', () => {
10+
circularA()
11+
12+
expect(circularB).toHaveBeenCalledOnce()
13+
})

test/core/test/mocked-process.test.ts renamed to test/core/test/stubbed-process.test.ts

File renamed without changes.

0 commit comments

Comments
 (0)