fix(vitest): throw an error if vi.mock is exported (#5034)

vitest-dev · Jan 23, 2024 · 253df1c · 253df1c
1 parent 7344870
commit 253df1c
Show file tree

Hide file tree

Showing 4 changed files with 104 additions and 38 deletions.
diff --git a/packages/vite-node/src/source-map-handler.ts b/packages/vite-node/src/source-map-handler.ts
@@ -11,14 +11,11 @@ import { TraceMap, originalPositionFor } from '@jridgewell/trace-mapping'
 // Only install once if called multiple times
 let errorFormatterInstalled = false
 
-// If true, the caches are reset before a stack trace formatting operation
-const emptyCacheBetweenOperations = false
-
 // Maps a file path to a string containing the file contents
-let fileContentsCache: Record<string, string> = {}
+const fileContentsCache: Record<string, string> = {}
 
 // Maps a file path to a source map for that file
-let sourceMapCache: Record<string, { url: string | null; map: TraceMap | null }> = {}
+const sourceMapCache: Record<string, { url: string | null; map: TraceMap | null }> = {}
 
 // Regex for detecting source maps
 const reSourceMap = /^data:application\/json[^,]+base64,/
@@ -405,11 +402,6 @@ function wrapCallSite(frame: CallSite, state: State) {
 // This function is part of the V8 stack trace API, for more info see:
 // https://v8.dev/docs/stack-trace-api
 function prepareStackTrace(error: Error, stack: CallSite[]) {
-  if (emptyCacheBetweenOperations) {
-    fileContentsCache = {}
-    sourceMapCache = {}
-  }
-
   const name = error.name || 'Error'
   const message = error.message || ''
   const errorString = `${name}: ${message}`

diff --git a/packages/vitest/src/node/hoistMocks.ts b/packages/vitest/src/node/hoistMocks.ts
@@ -162,6 +162,25 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
     }
   }
 
+  function assertNotDefaultExport(node: Positioned<CallExpression>, error: string) {
+    const defaultExport = findNodeAround(ast, node.start, 'ExportDefaultDeclaration')?.node as Positioned<ExportDefaultDeclaration> | undefined
+    if (defaultExport?.declaration === node || (defaultExport?.declaration.type === 'AwaitExpression' && defaultExport.declaration.argument === node))
+      throw createSyntaxError(defaultExport, error)
+  }
+
+  function assertNotNamedExport(node: Positioned<VariableDeclaration>, error: string) {
+    const nodeExported = findNodeAround(ast, node.start, 'ExportNamedDeclaration')?.node as Positioned<ExportNamedDeclaration> | undefined
+    if (nodeExported?.declaration === node)
+      throw createSyntaxError(nodeExported, error)
+  }
+
+  function getVariableDeclaration(node: Positioned<CallExpression>) {
+    const declarationNode = findNodeAround(ast, node.start, 'VariableDeclaration')?.node as Positioned<VariableDeclaration> | undefined
+    const init = declarationNode?.declarations[0]?.init
+    if (init && (init === node || (init.type === 'AwaitExpression' && init.argument === node)))
+      return declarationNode
+  }
+
   esmWalker(ast, {
     onIdentifier(id, info, parentStack) {
       const binding = idToImportMap.get(id.name)
@@ -197,38 +216,21 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
       ) {
         const methodName = node.callee.property.name
 
-        if (methodName === 'mock' || methodName === 'unmock')
+        if (methodName === 'mock' || methodName === 'unmock') {
+          const method = `${node.callee.object.name}.${methodName}`
+          assertNotDefaultExport(node, `Cannot export the result of "${method}". Remove export declaration because "${method}" doesn\'t return anything.`)
+          const declarationNode = getVariableDeclaration(node)
+          if (declarationNode)
+            assertNotNamedExport(declarationNode, `Cannot export the result of "${method}". Remove export declaration because "${method}" doesn\'t return anything.`)
           hoistedNodes.push(node)
+        }
 
         if (methodName === 'hoisted') {
-          // check it's not a default export
-          const defaultExport = findNodeAround(ast, node.start, 'ExportDefaultDeclaration')?.node as Positioned<ExportDefaultDeclaration> | undefined
-          if (defaultExport?.declaration === node || (defaultExport?.declaration.type === 'AwaitExpression' && defaultExport.declaration.argument === node))
-            throw createSyntaxError(defaultExport, 'Cannot export hoisted variable. You can control hoisting behavior by placing the import from this file first.')
-
-          const declarationNode = findNodeAround(ast, node.start, 'VariableDeclaration')?.node as Positioned<VariableDeclaration> | undefined
-          const init = declarationNode?.declarations[0]?.init
-          const isViHoisted = (node: CallExpression) => {
-            return node.callee.type === 'MemberExpression'
-              && isIdentifier(node.callee.object)
-              && (node.callee.object.name === 'vi' || node.callee.object.name === 'vitest')
-              && isIdentifier(node.callee.property)
-              && node.callee.property.name === 'hoisted'
-          }
+          assertNotDefaultExport(node, 'Cannot export hoisted variable. You can control hoisting behavior by placing the import from this file first.')
 
-          const canMoveDeclaration = (init
-            && init.type === 'CallExpression'
-            && isViHoisted(init)) /* const v = vi.hoisted() */
-            || (init
-                && init.type === 'AwaitExpression'
-                && init.argument.type === 'CallExpression'
-                && isViHoisted(init.argument)) /* const v = await vi.hoisted() */
-
-          if (canMoveDeclaration) {
-            // export const variable = vi.hoisted()
-            const nodeExported = findNodeAround(ast, declarationNode.start, 'ExportNamedDeclaration')?.node as Positioned<ExportNamedDeclaration> | undefined
-            if (nodeExported?.declaration === declarationNode)
-              throw createSyntaxError(nodeExported, 'Cannot export hoisted variable. You can control hoisting behavior by placing the import from this file first.')
+          const declarationNode = getVariableDeclaration(node)
+          if (declarationNode) {
+            assertNotNamedExport(declarationNode, 'Cannot export hoisted variable. You can control hoisting behavior by placing the import from this file first.')
             // hoist "const variable = vi.hoisted(() => {})"
             hoistedNodes.push(declarationNode)
           }

diff --git a/test/core/test/__snapshots__/injector-mock.test.ts.snap b/test/core/test/__snapshots__/injector-mock.test.ts.snap
@@ -120,3 +120,43 @@ exports[`throws an error when nodes are incompatible > correctly throws an error
       5| })
       6| "
 `;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.mock is exported as a named export 1`] = `"Cannot export the result of "vi.mock". Remove export declaration because "vi.mock" doesn't return anything."`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.mock is exported as a named export 2`] = `
+"      1| import { vi } from 'vitest'
+      2| 
+      3| export const mocked = vi.mock('./mocked')
+       | ^
+      4| "
+`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.mock is exported as default export 1`] = `"Cannot export the result of "vi.mock". Remove export declaration because "vi.mock" doesn't return anything."`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.mock is exported as default export 2`] = `
+"      1| import { vi } from 'vitest'
+      2| 
+      3| export default vi.mock('./mocked')
+       | ^
+      4| "
+`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.unmock is exported as a named export 1`] = `"Cannot export the result of "vi.unmock". Remove export declaration because "vi.unmock" doesn't return anything."`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.unmock is exported as a named export 2`] = `
+"      1| import { vi } from 'vitest'
+      2| 
+      3| export const mocked = vi.unmock('./mocked')
+       | ^
+      4| "
+`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.unmock is exported as default export 1`] = `"Cannot export the result of "vi.unmock". Remove export declaration because "vi.unmock" doesn't return anything."`;
+
+exports[`throws an error when nodes are incompatible > correctly throws an error if vi.unmock is exported as default export 2`] = `
+"      1| import { vi } from 'vitest'
+      2| 
+      3| export default vi.unmock('./mocked')
+       | ^
+      4| "
+`;
diff --git a/test/core/test/injector-mock.test.ts b/test/core/test/injector-mock.test.ts
@@ -1323,6 +1323,38 @@ import { vi } from 'vitest'
 export default await vi.hoisted(async () => {
   return {}
 })
+`,
+    ],
+    [
+      'vi.mock is exported as default export',
+      `\
+import { vi } from 'vitest'
+
+export default vi.mock('./mocked')
+`,
+    ],
+    [
+      'vi.unmock is exported as default export',
+      `\
+import { vi } from 'vitest'
+
+export default vi.unmock('./mocked')
+`,
+    ],
+    [
+      'vi.mock is exported as a named export',
+      `\
+import { vi } from 'vitest'
+
+export const mocked = vi.mock('./mocked')
+`,
+    ],
+    [
+      'vi.unmock is exported as a named export',
+      `\
+import { vi } from 'vitest'
+
+export const mocked = vi.unmock('./mocked')
 `,
     ],
   ])('correctly throws an error if %s', (_, code) => {