feat: allow import statement as vi.mock path for better IDE support (#5690)

sheremet-va · web-flow · commit a99a14c1cc7c · 2024-05-09T13:19:38.000+02:00
diff --git 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**: `<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`.
 
@@ -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
@@ -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.
@@ -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.
@@ -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,
@@ -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,
@@ -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())
     },
 
diff --git 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<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.')
 
@@ -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]
@@ -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('')
diff --git a/packages/vitest/src/types/mocker.ts b/packages/vitest/src/types/mocker.ts
@@ -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>>
diff --git 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
@@ -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/stubbed-process.test.ts b/test/core/test/stubbed-process.test.ts

Original file line number	Diff line number	Diff 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`
`2`	`2`	`export type MockFactory = () => any`
`3`	`3`
`4`	`4`	`export type MockMap = Map<string, Record<string, string \| null \| MockFactory>>`