Skip to content

Commit

Permalink
feat: allow import statement as vi.mock path for better IDE support (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
sheremet-va authored May 9, 2024
1 parent 8c96607 commit a99a14c
Show file tree
Hide file tree
Showing 7 changed files with 217 additions and 14 deletions.
16 changes: 16 additions & 0 deletions docs/api/vi.md
Original file line number Diff line number Diff line change
Expand Up @@ -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**: `<T>(path: Promise<T>, factory?: (importOriginal: () => T) => unknown) => void` <Version>2.0.0+</Version>

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

Expand Down Expand Up @@ -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.

Expand Down
36 changes: 28 additions & 8 deletions packages/vitest/src/integrations/vi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -184,15 +184,21 @@ 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<T>(module: Promise<T>, factory?: MockFactoryWithHelper<T>): void

/**
* Removes module from mocked registry. All calls to import will return the original module even if it was mocked before.
*
* 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<unknown>): void

/**
* Mocks every subsequent [dynamic import](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/import) call.
Expand All @@ -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<T>(module: Promise<T>, factory?: MockFactoryWithHelper<T>): 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<unknown>): void

/**
* Imports module, bypassing all checks if it should be mocked.
Expand Down Expand Up @@ -476,7 +488,9 @@ function createVitest(): VitestUtils {
return factory()
},

mock(path: string, factory?: MockFactoryWithHelper) {
mock(path: string | Promise<unknown>, 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,
Expand All @@ -486,11 +500,15 @@ function createVitest(): VitestUtils {
)
},

unmock(path: string) {
unmock(path: string | Promise<unknown>) {
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<unknown>, 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,
Expand All @@ -500,7 +518,9 @@ function createVitest(): VitestUtils {
)
},

doUnmock(path: string) {
doUnmock(path: string | Promise<unknown>) {
if (typeof path !== 'string')
throw new Error(`vi.doUnmock() expects a string path, but received a ${typeof path}`)
_mocker.queueUnmock(path, getImporter())
},

Expand Down
66 changes: 61 additions & 5 deletions packages/vitest/src/node/hoistMocks.ts
Original file line number Diff line number Diff line change
@@ -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'
Expand Down Expand Up @@ -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) {
Expand All @@ -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)
}
Expand Down Expand Up @@ -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<Expression>
let source: Positioned<Expression> | null = null
if (moduleInfo.type === 'ImportExpression')
source = moduleInfo.source as Positioned<Expression>
if (moduleInfo.type === 'AwaitExpression' && moduleInfo.argument.type === 'ImportExpression')
source = moduleInfo.argument.source as Positioned<Expression>
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.')

Expand Down Expand Up @@ -277,6 +295,14 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
)
}

function rewriteMockDynamicImport(nodeCode: string, moduleInfo: Positioned<ImportExpression>, expressionStart: number, expressionEnd: number, mockStart: number) {
const source = moduleInfo.source as Positioned<Expression>
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]
Expand All @@ -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<Expression>
// 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<ImportExpression>,
moduleInfo.start,
moduleInfo.end,
node.start,
)
}
}

s.remove(node.start, end)
return `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
}).join('')
Expand Down
2 changes: 1 addition & 1 deletion packages/vitest/src/types/mocker.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
export type MockFactoryWithHelper = (importOriginal: <T = unknown>() => Promise<T>) => any
export type MockFactoryWithHelper<M = unknown> = (importOriginal: <T extends M>() => Promise<T>) => any
export type MockFactory = () => any

export type MockMap = Map<string, Record<string, string | null | MockFactory>>
Expand Down
98 changes: 98 additions & 0 deletions test/core/test/injector-mock.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
13 changes: 13 additions & 0 deletions test/core/test/mocked-import-circular.test.ts
Original file line number Diff line number Diff line change
@@ -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()
})
File renamed without changes.

0 comments on commit a99a14c

Please sign in to comment.