Skip to content

Commit

Permalink
fix(browser): improve source maps when vi.mock is present (#6810)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sheremet-va
sheremet-va authored Nov 13, 2024
1 parent e26e066 commit 8d179af
Show file tree
Hide file tree
Showing 4 changed files with 331 additions and 266 deletions.
228 changes: 120 additions & 108 deletions packages/mocker/src/node/hoistMocksPlugin.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,23 +6,23 @@ import type {
Expression,
Identifier,
ImportDeclaration,
ImportExpression,
VariableDeclaration,
} from 'estree'
import type { SourceMap } from 'magic-string'
import type { RollupAstNode } from 'rollup'
import type { Plugin, Rollup } from 'vite'
import type { Node, Positioned } from './esmWalker'
import { findNodeAround } from 'acorn-walk'
import MagicString from 'magic-string'
import { createFilter } from 'vite'
import { esmWalker, getArbitraryModuleIdentifier } from './esmWalker'
import { esmWalker } from './esmWalker'

interface HoistMocksOptions {
/**
* List of modules that should always be imported before compiler hints.
* @default ['vitest']
* @default 'vitest'
*/
hoistedModules?: string[]
hoistedModule?: string
/**
* @default ["vi", "vitest"]
*/
Expand Down Expand Up @@ -106,11 +106,14 @@ function isIdentifier(node: any): node is Positioned<Identifier> {
return node.type === 'Identifier'
}

function getBetterEnd(code: string, node: Node) {
function getNodeTail(code: string, node: Node) {
let end = node.end
if (code[node.end] === ';') {
end += 1
}
if (code[node.end] === '\n') {
return end + 1
}
if (code[node.end + 1] === '\n') {
end += 1
}
Expand Down Expand Up @@ -160,48 +163,43 @@ export function hoistMocks(
dynamicImportMockMethodNames = ['mock', 'unmock', 'doMock', 'doUnmock'],
hoistedMethodNames = ['hoisted'],
utilsObjectNames = ['vi', 'vitest'],
hoistedModules = ['vitest'],
hoistedModule = 'vitest',
} = options

const hoistIndex = code.match(hashbangRE)?.[0].length ?? 0
// hoist at the start of the file, after the hashbang
let hoistIndex = hashbangRE.exec(code)?.[0].length ?? 0

let hoistedModuleImported = false

let uid = 0
const idToImportMap = new Map<string, string>()

const imports: {
node: RollupAstNode<ImportDeclaration>
id: string
}[] = []

// this will transform import statements into dynamic ones, if there are imports
// it will keep the import as is, if we don't need to mock anything
// in browser environment it will wrap the module value with "vitest_wrap_module" function
// that returns a proxy to the module so that named exports can be mocked
const transformImportDeclaration = (node: ImportDeclaration) => {
const source = node.source.value as string

const importId = `__vi_import_${uid++}__`
const hasSpecifiers = node.specifiers.length > 0
const code = hasSpecifiers
? `const ${importId} = await import('${source}')\n`
: `await import('${source}')\n`
return {
code,
id: importId,
}
}

function defineImport(node: Positioned<ImportDeclaration>) {
function defineImport(
importNode: ImportDeclaration & {
start: number
end: number
},
) {
const source = importNode.source.value as string
// always hoist vitest import to top of the file, so
// "vi" helpers can access it
if (hoistedModules.includes(node.source.value as string)) {
if (hoistedModule === source) {
hoistedModuleImported = true
return
}
const importId = `__vi_import_${uid++}__`
imports.push({ id: importId, node: importNode })

const declaration = transformImportDeclaration(node)
if (!declaration) {
return null
}
s.appendLeft(hoistIndex, declaration.code)
return declaration.id
return importId
}

// 1. check all import statements and record id -> importName map
Expand All @@ -214,13 +212,20 @@ export function hoistMocks(
if (!importId) {
continue
}
s.remove(node.start, getBetterEnd(code, node))
for (const spec of node.specifiers) {
if (spec.type === 'ImportSpecifier') {
idToImportMap.set(
spec.local.name,
`${importId}.${getArbitraryModuleIdentifier(spec.imported)}`,
)
if (spec.imported.type === 'Identifier') {
idToImportMap.set(
spec.local.name,
`${importId}.${spec.imported.name}`,
)
}
else {
idToImportMap.set(
spec.local.name,
`${importId}[${JSON.stringify(spec.imported.value as string)}]`,
)
}
}
else if (spec.type === 'ImportDefaultSpecifier') {
idToImportMap.set(spec.local.name, `${importId}.default`)
Expand All @@ -235,7 +240,7 @@ export function hoistMocks(

const declaredConst = new Set<string>()
const hoistedNodes: Positioned<
CallExpression | VariableDeclaration | AwaitExpression
CallExpression | VariableDeclaration | AwaitExpression
>[] = []

function createSyntaxError(node: Positioned<Node>, message: string) {
Expand Down Expand Up @@ -300,6 +305,8 @@ export function hoistMocks(
}
}

const usedUtilityExports = new Set<string>()

esmWalker(ast, {
onIdentifier(id, info, parentStack) {
const binding = idToImportMap.get(id.name)
Expand Down Expand Up @@ -333,6 +340,7 @@ export function hoistMocks(
&& isIdentifier(node.callee.property)
) {
const methodName = node.callee.property.name
usedUtilityExports.add(node.callee.object.name)

if (hoistableMockMethodNames.includes(methodName)) {
const method = `${node.callee.object.name}.${methodName}`
Expand All @@ -347,6 +355,35 @@ export function hoistMocks(
`Cannot export the result of "${method}". Remove export declaration because "${method}" doesn\'t return anything.`,
)
}
// rewrite vi.mock(import('..')) into vi.mock('..')
if (
node.type === 'CallExpression'
&& node.callee.type === 'MemberExpression'
&& dynamicImportMockMethodNames.includes((node.callee.property as Identifier).name)
) {
const moduleInfo = node.arguments[0] as Positioned<Expression>
// vi.mock(import('./path')) -> vi.mock('./path')
if (moduleInfo.type === 'ImportExpression') {
const source = moduleInfo.source as Positioned<Expression>
s.overwrite(
moduleInfo.start,
moduleInfo.end,
s.slice(source.start, source.end),
)
}
// vi.mock(await import('./path')) -> vi.mock('./path')
if (
moduleInfo.type === 'AwaitExpression'
&& moduleInfo.argument.type === 'ImportExpression'
) {
const source = moduleInfo.argument.source as Positioned<Expression>
s.overwrite(
moduleInfo.start,
moduleInfo.end,
s.slice(source.start, source.end),
)
}
}
hoistedNodes.push(node)
}
// vi.doMock(import('./path')) -> vi.doMock('./path')
Expand Down Expand Up @@ -394,9 +431,8 @@ export function hoistMocks(
'AwaitExpression',
)?.node as Positioned<AwaitExpression> | undefined
// hoist "await vi.hoisted(async () => {})" or "vi.hoisted(() => {})"
hoistedNodes.push(
awaitedExpression?.argument === node ? awaitedExpression : node,
)
const moveNode = awaitedExpression?.argument === node ? awaitedExpression : node
hoistedNodes.push(moveNode)
}
}
}
Expand Down Expand Up @@ -446,24 +482,6 @@ export function hoistMocks(
)
}

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 @@ -479,61 +497,55 @@ export function hoistMocks(
}
}

// Wait for imports to be hoisted and then hoist the mocks
const hoistedCode = hoistedNodes
.map((node) => {
const end = getBetterEnd(code, node)
/**
* In the following case, we need to change the `user` to user: __vi_import_x__.user
* So we should get the latest code from `s`.
*
* import user from './user'
* vi.mock('./mock.js', () => ({ getSession: vi.fn().mockImplementation(() => ({ user })) }))
*/
let nodeCode = s.slice(node.start, end)

// rewrite vi.mock(import('..')) into vi.mock('..')
if (
node.type === 'CallExpression'
&& node.callee.type === 'MemberExpression'
&& dynamicImportMockMethodNames.includes((node.callee.property as Identifier).name)
) {
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,
)
}
}
// hoist vi.mock/vi.hoisted
for (const node of hoistedNodes) {
const end = getNodeTail(code, node)
if (hoistIndex === end) {
hoistIndex = end
}
// don't hoist into itself if it's already at the top
else if (hoistIndex !== node.start) {
s.move(node.start, end, hoistIndex)
}
}

s.remove(node.start, end)
return `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
})
.join('')
// hoist actual dynamic imports last so they are inserted after all hoisted mocks
for (const { node: importNode, id: importId } of imports) {
const source = importNode.source.value as string

if (hoistedCode || hoistedModuleImported) {
s.prepend(
(!hoistedModuleImported && hoistedCode ? API_NOT_FOUND_CHECK(utilsObjectNames) : '')
+ hoistedCode,
s.update(
importNode.start,
importNode.end,
`const ${importId} = await import(${JSON.stringify(
source,
)});\n`,
)

if (importNode.start === hoistIndex) {
// no need to hoist, but update hoistIndex to keep the order
hoistIndex = importNode.end
}
else {
// There will be an error if the module is called before it is imported,
// so the module import statement is hoisted to the top
s.move(importNode.start, importNode.end, hoistIndex)
}
}

if (!hoistedModuleImported && hoistedNodes.length) {
const utilityImports = [...usedUtilityExports]
// "vi" or "vitest" is imported from a module other than "vitest"
if (utilityImports.some(name => idToImportMap.has(name))) {
s.prepend(API_NOT_FOUND_CHECK(utilityImports))
}
// if "vi" or "vitest" are not imported at all, import them
else if (utilityImports.length) {
s.prepend(
`import { ${[...usedUtilityExports].join(', ')} } from ${JSON.stringify(
hoistedModule,
)}\n`,
)
}
}

return {
Expand Down
Loading

0 comments on commit 8d179af

Please sign in to comment.