From a99a14c1cc7c16ff63c9f5f3b027443fecd4d8df Mon Sep 17 00:00:00 2001 From: Vladimir Date: Thu, 9 May 2024 13:19:38 +0200 Subject: [PATCH] feat: allow import statement as vi.mock path for better IDE support (#5690) --- docs/api/vi.md | 16 +++ packages/vitest/src/integrations/vi.ts | 36 +++++-- packages/vitest/src/node/hoistMocks.ts | 66 ++++++++++++- packages/vitest/src/types/mocker.ts | 2 +- test/core/test/injector-mock.test.ts | 98 +++++++++++++++++++ test/core/test/mocked-import-circular.test.ts | 13 +++ ...rocess.test.ts => stubbed-process.test.ts} | 0 7 files changed, 217 insertions(+), 14 deletions(-) create mode 100644 test/core/test/mocked-import-circular.test.ts rename test/core/test/{mocked-process.test.ts => stubbed-process.test.ts} (100%) diff --git a/docs/api/vi.md b/docs/api/vi.md index 9118ce845df2..17dc0d6dd35a 100644 --- a/docs/api/vi.md +++ b/docs/api/vi.md @@ -17,6 +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**: `(path: Promise, factory?: (importOriginal: () => T) => unknown) => void` 2.0.0+ 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`. @@ -64,6 +65,21 @@ vi.mock('./path/to/module.js', async (importOriginal) => { }) ``` +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). + +```ts +vi.mock(import('./path/to/module.js'), async (importOriginal) => { + const mod = await importOriginal() // type is inferred + return { + ...mod, + // replace some exports + namedExport: vi.fn(), + } +}) +``` + +Under the hood, Vitest still operates on a string and not a module object. + ::: warning `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. diff --git a/packages/vitest/src/integrations/vi.ts b/packages/vitest/src/integrations/vi.ts index 478bcff9ff70..64d5d2581304 100644 --- a/packages/vitest/src/integrations/vi.ts +++ b/packages/vitest/src/integrations/vi.ts @@ -184,7 +184,10 @@ export interface VitestUtils { * @param path Path to the module. Can be aliased, if your Vitest config supports it * @param factory Mocked module factory. The result of this function will be an exports object */ - mock: (path: string, factory?: MockFactoryWithHelper) => void + // eslint-disable-next-line ts/method-signature-style + mock(path: string, factory?: MockFactoryWithHelper): void + // eslint-disable-next-line ts/method-signature-style + mock(module: Promise, factory?: MockFactoryWithHelper): void /** * Removes module from mocked registry. All calls to import will return the original module even if it was mocked before. @@ -192,7 +195,10 @@ export interface VitestUtils { * This call is hoisted to the top of the file, so it will only unmock modules that were defined in `setupFiles`, for example. * @param path Path to the module. Can be aliased, if your Vitest config supports it */ - unmock: (path: string) => void + // eslint-disable-next-line ts/method-signature-style + unmock(path: string): void + // eslint-disable-next-line ts/method-signature-style + unmock(module: Promise): void /** * 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 { * @param path Path to the module. Can be aliased, if your Vitest config supports it * @param factory Mocked module factory. The result of this function will be an exports object */ - doMock: (path: string, factory?: MockFactoryWithHelper) => void + // eslint-disable-next-line ts/method-signature-style + doMock(path: string, factory?: MockFactoryWithHelper): void + // eslint-disable-next-line ts/method-signature-style + doMock(module: Promise, factory?: MockFactoryWithHelper): void /** * Removes module from mocked registry. All subsequent calls to import will return original module. * * Unlike [`vi.unmock`](https://vitest.dev/api/vi#vi-unmock), this method is not hoisted to the top of the file. * @param path Path to the module. Can be aliased, if your Vitest config supports it */ - doUnmock: (path: string) => void + // eslint-disable-next-line ts/method-signature-style + doUnmock(path: string): void + // eslint-disable-next-line ts/method-signature-style + doUnmock(module: Promise): void /** * Imports module, bypassing all checks if it should be mocked. @@ -476,7 +488,9 @@ function createVitest(): VitestUtils { return factory() }, - mock(path: string, factory?: MockFactoryWithHelper) { + mock(path: string | Promise, factory?: MockFactoryWithHelper) { + if (typeof path !== 'string') + throw new Error(`vi.mock() expects a string path, but received a ${typeof path}`) const importer = getImporter() _mocker.queueMock( path, @@ -486,11 +500,15 @@ function createVitest(): VitestUtils { ) }, - unmock(path: string) { + unmock(path: string | Promise) { + if (typeof path !== 'string') + throw new Error(`vi.unmock() expects a string path, but received a ${typeof path}`) _mocker.queueUnmock(path, getImporter()) }, - doMock(path: string, factory?: MockFactoryWithHelper) { + doMock(path: string | Promise, factory?: MockFactoryWithHelper) { + if (typeof path !== 'string') + throw new Error(`vi.doMock() expects a string path, but received a ${typeof path}`) const importer = getImporter() _mocker.queueMock( path, @@ -500,7 +518,9 @@ function createVitest(): VitestUtils { ) }, - doUnmock(path: string) { + doUnmock(path: string | Promise) { + if (typeof path !== 'string') + throw new Error(`vi.doUnmock() expects a string path, but received a ${typeof path}`) _mocker.queueUnmock(path, getImporter()) }, diff --git a/packages/vitest/src/node/hoistMocks.ts b/packages/vitest/src/node/hoistMocks.ts index bc832cee22f5..bb05aca6967f 100644 --- a/packages/vitest/src/node/hoistMocks.ts +++ b/packages/vitest/src/node/hoistMocks.ts @@ -1,7 +1,7 @@ import MagicString from 'magic-string' -import type { AwaitExpression, CallExpression, ExportDefaultDeclaration, ExportNamedDeclaration, Identifier, ImportDeclaration, VariableDeclaration, Node as _Node } from 'estree' +import type { AwaitExpression, CallExpression, ExportDefaultDeclaration, ExportNamedDeclaration, Expression, Identifier, ImportDeclaration, ImportExpression, VariableDeclaration, Node as _Node } from 'estree' import { findNodeAround } from 'acorn-walk' -import type { PluginContext } from 'rollup' +import type { PluginContext, ProgramNode } from 'rollup' import { esmWalker } from '@vitest/utils/ast' import type { Colors } from '@vitest/utils' import { highlightCode } from '../utils/colors' @@ -58,7 +58,7 @@ export function getBetterEnd(code: string, node: Node) { return end } -const regexpHoistable = /\b(vi|vitest)\s*\.\s*(mock|unmock|hoisted)\(/ +const regexpHoistable = /\b(vi|vitest)\s*\.\s*(mock|unmock|hoisted|doMock|doUnmock)\(/ const hashbangRE = /^#!.*\n/ 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 const s = new MagicString(code) - let ast: any + let ast: ProgramNode try { ast = parse(code) } @@ -225,6 +225,24 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse hoistedNodes.push(node) } + // vi.doMock(import('./path')) -> vi.doMock('./path') + // vi.doMock(await import('./path')) -> vi.doMock('./path') + if (methodName === 'doMock' || methodName === 'doUnmock') { + const moduleInfo = node.arguments[0] as Positioned + let source: Positioned | null = null + if (moduleInfo.type === 'ImportExpression') + source = moduleInfo.source as Positioned + if (moduleInfo.type === 'AwaitExpression' && moduleInfo.argument.type === 'ImportExpression') + source = moduleInfo.argument.source as Positioned + if (source) { + s.overwrite( + moduleInfo.start, + moduleInfo.end, + s.slice(source.start, source.end), + ) + } + } + if (methodName === 'hoisted') { assertNotDefaultExport(node, 'Cannot export hoisted variable. You can control hoisting behavior by placing the import from this file first.') @@ -277,6 +295,14 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse ) } + function rewriteMockDynamicImport(nodeCode: string, moduleInfo: Positioned, expressionStart: number, expressionEnd: number, mockStart: number) { + const source = moduleInfo.source as Positioned + const importPath = s.slice(source.start, source.end) + const nodeCodeStart = expressionStart - mockStart + const nodeCodeEnd = expressionEnd - mockStart + return nodeCode.slice(0, nodeCodeStart) + importPath + nodeCode.slice(nodeCodeEnd) + } + // validate hoistedNodes doesn't have nodes inside other nodes for (let i = 0; i < hoistedNodes.length; i++) { const node = hoistedNodes[i] @@ -300,7 +326,37 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse * import user from './user' * vi.mock('./mock.js', () => ({ getSession: vi.fn().mockImplementation(() => ({ user })) })) */ - const nodeCode = s.slice(node.start, end) + let nodeCode = s.slice(node.start, end) + + // rewrite vi.mock(import('..')) into vi.mock('..') + if ( + node.type === 'CallExpression' + && node.callee.type === 'MemberExpression' + && ((node.callee.property as Identifier).name === 'mock' || (node.callee.property as Identifier).name === 'unmock') + ) { + const moduleInfo = node.arguments[0] as Positioned + // vi.mock(import('./path')) -> vi.mock('./path') + if (moduleInfo.type === 'ImportExpression') { + nodeCode = rewriteMockDynamicImport( + nodeCode, + moduleInfo, + moduleInfo.start, + moduleInfo.end, + node.start, + ) + } + // vi.mock(await import('./path')) -> vi.mock('./path') + if (moduleInfo.type === 'AwaitExpression' && moduleInfo.argument.type === 'ImportExpression') { + nodeCode = rewriteMockDynamicImport( + nodeCode, + moduleInfo.argument as Positioned, + moduleInfo.start, + moduleInfo.end, + node.start, + ) + } + } + s.remove(node.start, end) return `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}` }).join('') diff --git a/packages/vitest/src/types/mocker.ts b/packages/vitest/src/types/mocker.ts index 76f759de3318..efd99409be1d 100644 --- a/packages/vitest/src/types/mocker.ts +++ b/packages/vitest/src/types/mocker.ts @@ -1,4 +1,4 @@ -export type MockFactoryWithHelper = (importOriginal: () => Promise) => any +export type MockFactoryWithHelper = (importOriginal: () => Promise) => any export type MockFactory = () => any export type MockMap = Map> diff --git a/test/core/test/injector-mock.test.ts b/test/core/test/injector-mock.test.ts index 047718175b33..ebd8d5ab77fa 100644 --- a/test/core/test/injector-mock.test.ts +++ b/test/core/test/injector-mock.test.ts @@ -1201,6 +1201,104 @@ await vi 1234;" `) }) + + test('handles dynamic import as the first argument', () => { + expect( + hoistSimpleCode(` + vi.mock(import('./path')) + vi.mock(import(somePath)) + vi.mock(import(\`./path\`)) + + vi.mock(import('./path')); + vi.mock(import(somePath)); + vi.mock(import(\`./path\`)); + + vi.mock(await import('./path')) + vi.mock(await import(somePath)) + vi.mock(await import(\`./path\`)) + + vi.mock(await import('./path')); + vi.mock(await import(somePath)); + vi.mock(await import(\`./path\`)); + + vi.mock(import('./path'), () => {}) + vi.mock(import(somePath), () => {}) + vi.mock(import(\`./path\`), () => {}) + + vi.mock(await import('./path'), () => {}) + vi.mock(await import(somePath), () => {}) + vi.mock(await import(\`./path\`), () => {}) + + vi.mock(import('./path'), () => {}); + vi.mock(import(somePath), () => {}); + vi.mock(import(\`./path\`), () => {}); + + vi.mock(await import('./path'), () => {}); + vi.mock(await import(somePath), () => {}); + vi.mock(await import(\`./path\`), () => {}); + `), + ).toMatchInlineSnapshot(` + "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") } + vi.mock('./path') + vi.mock(somePath) + vi.mock(\`./path\`) + vi.mock('./path'); + vi.mock(somePath); + vi.mock(\`./path\`); + vi.mock('./path') + vi.mock(somePath) + vi.mock(\`./path\`) + vi.mock('./path'); + vi.mock(somePath); + vi.mock(\`./path\`); + vi.mock('./path', () => {}) + vi.mock(somePath, () => {}) + vi.mock(\`./path\`, () => {}) + vi.mock('./path', () => {}) + vi.mock(somePath, () => {}) + vi.mock(\`./path\`, () => {}) + vi.mock('./path', () => {}); + vi.mock(somePath, () => {}); + vi.mock(\`./path\`, () => {}); + vi.mock('./path', () => {}); + vi.mock(somePath, () => {}); + vi.mock(\`./path\`, () => {});" + `) + }) + + test.only('handles import in vi.do* methods', () => { + expect( + hoistSimpleCode(` +vi.doMock(import('./path')) +vi.doMock(import(\`./path\`)) +vi.doMock(import('./path')); + +beforeEach(() => { + vi.doUnmock(import('./path')) + vi.doMock(import('./path')) +}) + +test('test', async () => { + vi.doMock(import(dynamicName)) + await import(dynamicName) +}) + `), + ).toMatchInlineSnapshot(` + "vi.doMock('./path') + vi.doMock(\`./path\`) + vi.doMock('./path'); + + beforeEach(() => { + vi.doUnmock('./path') + vi.doMock('./path') + }) + + test('test', async () => { + vi.doMock(dynamicName) + await import(dynamicName) + })" + `) + }) }) describe('throws an error when nodes are incompatible', () => { diff --git a/test/core/test/mocked-import-circular.test.ts b/test/core/test/mocked-import-circular.test.ts new file mode 100644 index 000000000000..98d766efba8a --- /dev/null +++ b/test/core/test/mocked-import-circular.test.ts @@ -0,0 +1,13 @@ +import { expect, it, vi } from 'vitest' + +// The order of the two imports here matters: B before A +import { circularB } from '../src/circularB' +import { circularA } from '../src/circularA' + +vi.mock(import('../src/circularB')) + +it('circular', () => { + circularA() + + expect(circularB).toHaveBeenCalledOnce() +}) diff --git a/test/core/test/mocked-process.test.ts b/test/core/test/stubbed-process.test.ts similarity index 100% rename from test/core/test/mocked-process.test.ts rename to test/core/test/stubbed-process.test.ts