feat: vi.mock can be called inside external libraries (#560)

Co-authored-by: Anthony Fu <anthonyfu117@hotmail.com>

examples/mocks/src/external.mjs

@@ -0,0 +1 @@
+vi.doMock('axios')

examples/mocks/test/external.test.ts

@@ -0,0 +1,7 @@
+import '../src/external.mjs'
+import { expect, test, vi } from 'vitest'
+import axios from 'axios'
+
+test('axios is mocked', () => {
+  expect(vi.isMockFunction(axios.get)).toBe(true)
+})

examples/mocks/vite.config.ts

@@ -7,5 +7,8 @@ export default defineConfig({
   test: {
     globals: true,
     environment: 'node',
+    deps: {
+      external: [/src\/external\.mjs/],
+    },
   },
 })

packages/vitest/src/integrations/vi.ts

@@ -1,6 +1,8 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 
 import mockdate from 'mockdate'
+import { parseStacktrace } from '../utils/source-map'
+import type { VitestMocker } from '../node/mocker'
 import { FakeTimers } from './timers'
 import type { EnhancedSpy, MaybeMocked, MaybeMockedDeep } from './jest-mock'
 import { fn, isMockFunction, spies, spyOn } from './jest-mock'
@@ -8,10 +10,16 @@ import { fn, isMockFunction, spies, spyOn } from './jest-mock'
 class VitestUtils {
   private _timers: FakeTimers
   private _mockedDate: string | number | Date | null
+  private _mocker: VitestMocker
 
   constructor() {
     this._timers = new FakeTimers()
+    // @ts-expect-error injected by vite-nide
+    this._mocker = typeof __vitest_mocker__ !== 'undefined' ? __vitest_mocker__ : null
     this._mockedDate = null
+
+    if (!this._mocker)
+      throw new Error('Vitest was initialised with native Node instead of Vite Node')
   }
 
   // timers
@@ -65,7 +73,11 @@ class VitestUtils {
   spyOn = spyOn
   fn = fn
 
-  // just hints for transformer to rewrite imports
+  private getImporter() {
+    const err = new Error('mock')
+    const [,, importer] = parseStacktrace(err, true)
+    return importer.file
+  }
 
   /**
    * Makes all `imports` to passed module to be mocked.
@@ -78,13 +90,26 @@ class VitestUtils {
    * @param path Path to the module. Can be aliased, if your config suppors it
    * @param factory Factory for the mocked module. Has the highest priority.
    */
-  public mock(path: string, factory?: () => any) {}
+  public mock(path: string, factory?: () => any) {
+    this._mocker.queueMock(path, this.getImporter(), factory)
+  }
+
   /**
    * Removes module from mocked registry. All subsequent calls to import will
    * return original module even if it was mocked.
    * @param path Path to the module. Can be aliased, if your config suppors it
    */
-  public unmock(path: string) {}
+  public unmock(path: string) {
+    this._mocker.queueUnmock(path, this.getImporter())
+  }
+
+  public doMock(path: string, factory?: () => any) {
+    this._mocker.queueMock(path, this.getImporter(), factory)
+  }
+
+  public doUnmock(path: string) {
+    this._mocker.queueUnmock(path, this.getImporter())
+  }
 
   /**
    * Imports module, bypassing all checks if it should be mocked.
@@ -99,7 +124,7 @@ class VitestUtils {
    * @returns Actual module without spies
    */
   public async importActual<T>(path: string): Promise<T> {
-    return {} as T
+    return this._mocker.importActual<T>(path, this.getImporter())
   }
 
   /**
@@ -109,7 +134,7 @@ class VitestUtils {
    * @returns Fully mocked module
    */
   public async importMock<T>(path: string): Promise<MaybeMockedDeep<T>> {
-    return {} as MaybeMockedDeep<T>
+    return this._mocker.importMock(path, this.getImporter())
   }
 
   /**
@@ -139,22 +164,19 @@ class VitestUtils {
   }
 
   public clearAllMocks() {
-    // @ts-expect-error clearing module mocks
-    __vitest__clearMocks__({ clearMocks: true })
+    this._mocker.clearMocks({ clearMocks: true })
     spies.forEach(spy => spy.mockClear())
     return this
   }
 
   public resetAllMocks() {
-    // @ts-expect-error resetting module mocks
-    __vitest__clearMocks__({ mockReset: true })
+    this._mocker.clearMocks({ mockReset: true })
     spies.forEach(spy => spy.mockReset())
     return this
   }
 
   public restoreAllMocks() {
-    // @ts-expect-error restoring module mocks
-    __vitest__clearMocks__({ restoreMocks: true })
+    this._mocker.clearMocks({ restoreMocks: true })
     spies.forEach(spy => spy.mockRestore())
     return this
   }

packages/vitest/src/node/execute.ts

@@ -1,8 +1,7 @@
 import { ViteNodeRunner } from 'vite-node/client'
-import { toFilePath } from 'vite-node/utils'
-import type { ViteNodeRunnerOptions } from 'vite-node'
+import type { ModuleCache, ViteNodeRunnerOptions } from 'vite-node'
 import type { SuiteMocks } from './mocker'
-import { createMocker } from './mocker'
+import { VitestMocker } from './mocker'
 
 export interface ExecuteOptions extends ViteNodeRunnerOptions {
   files: string[]
@@ -20,76 +19,27 @@ export async function executeInViteNode(options: ExecuteOptions) {
 }
 
 export class VitestRunner extends ViteNodeRunner {
-  mocker: ReturnType<typeof createMocker>
+  mocker: VitestMocker
 
   constructor(public options: ExecuteOptions) {
     super(options)
-    this.mocker = createMocker(this.root, options.mockMap)
+    this.mocker = new VitestMocker(options, this.moduleCache)
   }
 
   prepareContext(context: Record<string, any>) {
     const request = context.__vite_ssr_import__
 
-    const callFunctionMock = async(dep: string, mock: () => any) => {
-      const cacheName = `${dep}__mock`
-      const cached = this.moduleCache.get(cacheName)?.exports
-      if (cached)
-        return cached
-      const exports = await mock()
-      this.setCache(cacheName, { exports })
-      return exports
-    }
+    const mocker = this.mocker.withRequest(request)
 
-    const requestWithMock = async(dep: string) => {
-      const mock = this.mocker.getDependencyMock(dep)
-      if (mock === null) {
-        const cacheName = `${dep}__mock`
-        const cache = this.moduleCache.get(cacheName)
-        if (cache?.exports)
-          return cache.exports
-        const cacheKey = toFilePath(dep, this.root)
-        const mod = this.moduleCache.get(cacheKey)?.exports || await request(dep)
-        const exports = this.mocker.mockObject(mod)
-        this.setCache(cacheName, { exports })
-        return exports
-      }
-      if (typeof mock === 'function')
-        return callFunctionMock(dep, mock)
-      if (typeof mock === 'string')
-        dep = mock
-      return request(dep)
-    }
-    const importActual = (path: string, external: string | null) => {
-      return request(this.mocker.getActualPath(path, external))
-    }
-    const importMock = async(path: string, external: string | null): Promise<any> => {
-      let mock = this.mocker.getDependencyMock(path)
-
-      if (mock === undefined)
-        mock = this.mocker.resolveMockPath(path, this.root, external)
-
-      if (mock === null) {
-        const fsPath = this.mocker.getActualPath(path, external)
-        const mod = await request(fsPath)
-        return this.mocker.mockObject(mod)
-      }
-      if (typeof mock === 'function')
-        return callFunctionMock(path, mock)
-      return requestWithMock(mock)
-    }
+    mocker.on('mocked', (dep: string, module: Partial<ModuleCache>) => {
+      this.setCache(dep, module)
+    })
 
     return Object.assign(context, {
-      __vite_ssr_import__: requestWithMock,
-      __vite_ssr_dynamic_import__: requestWithMock,
+      __vite_ssr_import__: (dep: string) => mocker.requestWithMock(dep),
+      __vite_ssr_dynamic_import__: (dep: string) => mocker.requestWithMock(dep),
 
-      // vitest.mock API
-      __vitest__mock__: this.mocker.mockPath,
-      __vitest__unmock__: this.mocker.unmockPath,
-      __vitest__importActual__: importActual,
-      __vitest__importMock__: importMock,
-      // spies from 'jest-mock' are different inside suites and execute,
-      // so wee need to call this twice - inside suite and here
-      __vitest__clearMocks__: this.mocker.clearMocks,
+      __vitest_mocker__: mocker,
     })
   }
 }

packages/vitest/src/node/mocker.ts

@@ -1,38 +1,21 @@
 import { existsSync, readdirSync } from 'fs'
 import { isNodeBuiltin } from 'mlly'
 import { basename, dirname, join, resolve } from 'pathe'
+import type { ModuleCache } from 'vite-node'
+import { toFilePath } from 'vite-node/utils'
 import { spies, spyOn } from '../integrations/jest-mock'
 import { mergeSlashes, normalizeId } from '../utils'
+import type { ExecuteOptions } from './execute'
 
 export type SuiteMocks = Record<string, Record<string, string | null | (() => unknown)>>
 
-function resolveMockPath(mockPath: string, root: string, external: string | null) {
-  const path = normalizeId(external || mockPath)
+type Callback = (...args: any[]) => unknown
 
-  // it's a node_module alias
-  // all mocks should be inside <root>/__mocks__
-  if (external || isNodeBuiltin(mockPath)) {
-    const mockDirname = dirname(path) // for nested mocks: @vueuse/integration/useJwt
-    const baseFilename = basename(path)
-    const mockFolder = resolve(root, '__mocks__', mockDirname)
-
-    if (!existsSync(mockFolder)) return null
-
-    const files = readdirSync(mockFolder)
-
-    for (const file of files) {
-      const [basename] = file.split('.')
-      if (basename === baseFilename)
-        return resolve(mockFolder, file).replace(root, '')
-    }
-
-    return null
-  }
-
-  const dir = dirname(path)
-  const baseId = basename(path)
-  const fullPath = resolve(dir, '__mocks__', baseId)
-  return existsSync(fullPath) ? fullPath.replace(root, '') : null
+interface PendingSuiteMock {
+  id: string
+  importer: string
+  type: 'mock' | 'unmock'
+  factory?: () => unknown
 }
 
 function getObjectType(value: unknown): string {
@@ -57,41 +40,47 @@ function mockPrototype(proto: any) {
   return newProto
 }
 
-function mockObject(obj: any) {
-  const type = getObjectType(obj)
+const pendingIds: PendingSuiteMock[] = []
 
-  if (Array.isArray(obj))
-    return []
-  else if (type !== 'Object' && type !== 'Module')
-    return obj
+export class VitestMocker {
+  private request!: (dep: string) => unknown
 
-  const newObj = { ...obj }
+  private root: string
+  // private mockMap: SuiteMocks
 
-  const proto = mockPrototype(Object.getPrototypeOf(obj))
-  Object.setPrototypeOf(newObj, proto)
+  private callbacks: Record<string, ((...args: any[]) => unknown)[]> = {}
 
-  // eslint-disable-next-line no-restricted-syntax
-  for (const k in obj) {
-    newObj[k] = mockObject(obj[k])
-    const type = getObjectType(obj[k])
-
-    if (type.includes('Function') && !obj[k].__isSpy) {
-      spyOn(newObj, k).mockImplementation(() => {})
-      Object.defineProperty(newObj[k], 'length', { value: 0 }) // tinyspy retains length, but jest doesnt
-    }
+  constructor(
+    public options: ExecuteOptions,
+    private moduleCache: Map<string, ModuleCache>,
+    request?: (dep: string) => unknown,
+  ) {
+    this.root = this.options.root
+    // this.mockMap = options.mockMap
+    this.request = request!
   }
-  return newObj
-}
 
-export function createMocker(root: string, mockMap: SuiteMocks) {
-  function getSuiteFilepath() {
+  get mockMap() {
+    return this.options.mockMap
+  }
+
+  public on(event: string, cb: Callback) {
+    this.callbacks[event] ??= []
+    this.callbacks[event].push(cb)
+  }
+
+  private emit(event: string, ...args: any[]) {
+    (this.callbacks[event] ?? []).forEach(fn => fn(...args))
+  }
+
+  public getSuiteFilepath() {
     return process.__vitest_worker__?.filepath || 'global'
   }
 
-  function getMocks() {
-    const suite = getSuiteFilepath()
-    const suiteMocks = mockMap[suite || '']
-    const globalMocks = mockMap.global
+  public getMocks() {
+    const suite = this.getSuiteFilepath()
+    const suiteMocks = this.mockMap[suite || '']
+    const globalMocks = this.mockMap.global
 
     return {
       ...suiteMocks,
@@ -99,34 +88,177 @@ export function createMocker(root: string, mockMap: SuiteMocks) {
     }
   }
 
-  function getDependencyMock(dep: string) {
-    return getMocks()[resolveDependency(dep)]
+  private async resolvePath(id: string, importer: string) {
+    const path = await this.options.resolveId(id, importer)
+    return {
+      path: normalizeId(path?.id || id),
+      external: path?.id.includes('/node_modules/') ? id : null,
+    }
   }
 
-  function getActualPath(path: string, external: string | null) {
+  private async resolveMocks() {
+    await Promise.all(pendingIds.map(async(mock) => {
+      const { path, external } = await this.resolvePath(mock.id, mock.importer)
+      if (mock.type === 'unmock')
+        this.unmockPath(path, external)
+      if (mock.type === 'mock')
+        this.mockPath(path, external, mock.factory)
+    }))
+
+    pendingIds.length = 0
+  }
+
+  private async callFunctionMock(dep: string, mock: () => any) {
+    const cacheName = `${dep}__mock`
+    const cached = this.moduleCache.get(cacheName)?.exports
+    if (cached)
+      return cached
+    const exports = await mock()
+    this.emit('mocked', cacheName, { exports })
+    return exports
+  }
+
+  public getDependencyMock(dep: string) {
+    return this.getMocks()[this.resolveDependency(dep)]
+  }
+
+  // npm resolves as /node_modules, but we store as /@fs/.../node_modules
+  public resolveDependency(dep: string) {
+    if (dep.startsWith('/node_modules/'))
+      dep = mergeSlashes(`/@fs/${join(this.root, dep)}`)
+
+    return normalizeId(dep)
+  }
+
+  public getActualPath(path: string, external: string | null) {
     if (external)
       return mergeSlashes(`/@fs/${path}`)
 
-    return normalizeId(path.replace(root, ''))
+    return normalizeId(path.replace(this.root, ''))
   }
 
-  function unmockPath(path: string, external: string | null) {
-    const suitefile = getSuiteFilepath()
+  public resolveMockPath(mockPath: string, external: string | null) {
+    const path = normalizeId(external || mockPath)
 
-    const fsPath = getActualPath(path, external)
-    mockMap[suitefile] ??= {}
-    delete mockMap[suitefile][fsPath]
+    // it's a node_module alias
+    // all mocks should be inside <root>/__mocks__
+    if (external || isNodeBuiltin(mockPath)) {
+      const mockDirname = dirname(path) // for nested mocks: @vueuse/integration/useJwt
+      const baseFilename = basename(path)
+      const mockFolder = resolve(this.root, '__mocks__', mockDirname)
+
+      if (!existsSync(mockFolder)) return null
+
+      const files = readdirSync(mockFolder)
+
+      for (const file of files) {
+        const [basename] = file.split('.')
+        if (basename === baseFilename)
+          return resolve(mockFolder, file).replace(this.root, '')
+      }
+
+      return null
+    }
+
+    const dir = dirname(path)
+    const baseId = basename(path)
+    const fullPath = resolve(dir, '__mocks__', baseId)
+    return existsSync(fullPath) ? fullPath.replace(this.root, '') : null
   }
 
-  function mockPath(path: string, external: string | null, factory?: () => any) {
-    const suitefile = getSuiteFilepath()
+  public mockObject(obj: any) {
+    const type = getObjectType(obj)
 
-    const fsPath = getActualPath(path, external)
-    mockMap[suitefile] ??= {}
-    mockMap[suitefile][fsPath] = factory || resolveMockPath(path, root, external)
+    if (Array.isArray(obj))
+      return []
+    else if (type !== 'Object' && type !== 'Module')
+      return obj
+
+    const newObj = { ...obj }
+
+    const proto = mockPrototype(Object.getPrototypeOf(obj))
+    Object.setPrototypeOf(newObj, proto)
+
+    // eslint-disable-next-line no-restricted-syntax
+    for (const k in obj) {
+      newObj[k] = this.mockObject(obj[k])
+      const type = getObjectType(obj[k])
+
+      if (type.includes('Function') && !obj[k].__isSpy) {
+        spyOn(newObj, k).mockImplementation(() => {})
+        Object.defineProperty(newObj[k], 'length', { value: 0 }) // tinyspy retains length, but jest doesnt
+      }
+    }
+    return newObj
   }
 
-  function clearMocks({ clearMocks, mockReset, restoreMocks }: { clearMocks: boolean; mockReset: boolean; restoreMocks: boolean }) {
+  public unmockPath(path: string, external: string | null) {
+    const suitefile = this.getSuiteFilepath()
+
+    const fsPath = this.getActualPath(path, external)
+
+    if (this.mockMap[suitefile]?.[fsPath])
+      delete this.mockMap[suitefile][fsPath]
+  }
+
+  public mockPath(path: string, external: string | null, factory?: () => any) {
+    const suitefile = this.getSuiteFilepath()
+
+    const fsPath = this.getActualPath(path, external)
+
+    this.mockMap[suitefile] ??= {}
+    this.mockMap[suitefile][fsPath] = factory || this.resolveMockPath(path, external)
+  }
+
+  public async importActual<T>(id: string, importer: string): Promise<T> {
+    const { path, external } = await this.resolvePath(id, importer)
+    const fsPath = this.getActualPath(path, external)
+    const result = await this.request(fsPath)
+    return result as T
+  }
+
+  public async importMock(id: string, importer: string): Promise<any> {
+    const { path, external } = await this.resolvePath(id, importer)
+
+    let mock = this.getDependencyMock(path)
+
+    if (mock === undefined)
+      mock = this.resolveMockPath(path, external)
+
+    if (mock === null) {
+      const fsPath = this.getActualPath(path, external)
+      const mod = await this.request(fsPath)
+      return this.mockObject(mod)
+    }
+    if (typeof mock === 'function')
+      return this.callFunctionMock(path, mock)
+    return this.requestWithMock(mock)
+  }
+
+  public async requestWithMock(dep: string) {
+    await this.resolveMocks()
+
+    const mock = this.getDependencyMock(dep)
+
+    if (mock === null) {
+      const cacheName = `${dep}__mock`
+      const cache = this.moduleCache.get(cacheName)
+      if (cache?.exports)
+        return cache.exports
+      const cacheKey = toFilePath(dep, this.root)
+      const mod = this.moduleCache.get(cacheKey)?.exports || await this.request(dep)
+      const exports = this.mockObject(mod)
+      this.emit('mocked', cacheName, { exports })
+      return exports
+    }
+    if (typeof mock === 'function')
+      return this.callFunctionMock(dep, mock)
+    if (typeof mock === 'string')
+      dep = mock
+    return this.request(dep)
+  }
+
+  public clearMocks({ clearMocks, mockReset, restoreMocks }: { clearMocks?: boolean; mockReset?: boolean; restoreMocks?: boolean }) {
     if (!clearMocks && !mockReset && !restoreMocks)
       return
 
@@ -140,24 +272,15 @@ export function createMocker(root: string, mockMap: SuiteMocks) {
     })
   }
 
-  // npm resolves as /node_modules, but we store as /@fs/.../node_modules
-  function resolveDependency(dep: string) {
-    if (dep.startsWith('/node_modules/'))
-      return mergeSlashes(`/@fs/${join(root, dep)}`)
-
-    return normalizeId(dep)
+  public queueMock(id: string, importer: string, factory?: () => unknown) {
+    pendingIds.push({ type: 'mock', id, importer, factory })
   }
 
-  return {
-    mockPath,
-    unmockPath,
-    clearMocks,
-    getActualPath,
-    getMocks,
-    getDependencyMock,
+  public queueUnmock(id: string, importer: string) {
+    pendingIds.push({ type: 'unmock', id, importer })
+  }
 
-    mockObject,
-    getSuiteFilepath,
-    resolveMockPath,
+  public withRequest(request: (dep: string) => unknown) {
+    return new VitestMocker(this.options, this.moduleCache, request)
   }
 }

packages/vitest/src/node/plugins/mock.ts

@@ -2,58 +2,42 @@ import type { Plugin } from 'vite'
 import MagicString from 'magic-string'
 import { getCallLastIndex } from '../../utils'
 
-const mockRegexp = /^ *\b((?:vitest|vi)\s*.\s*mock\(["`'\s]+(.*[@\w_-]+)["`'\s]+)[),]{1};?/gm
-const pathRegexp = /\b(?:vitest|vi)\s*.\s*(unmock|importActual|importMock)\(["`'\s](.*[@\w_-]+)["`'\s]\);?/mg
+const hoistRegexp = /^ *\b((?:vitest|vi)\s*.\s*(mock|unmock)\(["`'\s]+(.*[@\w_-]+)["`'\s]+)[),]{1};?/gm
 const vitestRegexp = /import {[^}]*}.*(?=["'`]vitest["`']).*/gm
 
+export function hoistMocks(code: string) {
+  let m: MagicString | undefined
+  const mocks = code.matchAll(hoistRegexp)
+
+  for (const mockResult of mocks) {
+    const lastIndex = getMockLastIndex(code.slice(mockResult.index!))
+
+    if (lastIndex === null) continue
+
+    const startIndex = mockResult.index!
+
+    const { insideComment, insideString } = getIndexStatus(code, startIndex)
+
+    if (insideComment || insideString)
+      continue
+
+    const endIndex = startIndex + lastIndex
+
+    m ??= new MagicString(code)
+
+    m.prepend(`${m.slice(startIndex, endIndex)}\n`)
+    m.remove(startIndex, endIndex)
+  }
+
+  return m
+}
+
 export const MocksPlugin = (): Plugin => {
   return {
     name: 'vitest:mock-plugin',
     enforce: 'post',
-    async transform(code, id) {
-      let m: MagicString | undefined
-      const matchAll = code.matchAll(pathRegexp)
-
-      for (const match of matchAll) {
-        const [line, method, modulePath] = match
-        const filepath = await this.resolve(modulePath, id)
-        m ??= new MagicString(code)
-        const start = match.index || 0
-        const end = start + line.length
-
-        const overwrite = `${getMethodCall(method, filepath?.id || modulePath, modulePath)});`
-
-        m.overwrite(start, end, overwrite)
-      }
-
-      const mocks = code.matchAll(mockRegexp)
-
-      for (const mockResult of mocks) {
-        // we need to parse parsed string because factory may contain importActual
-        const lastIndex = getMockLastIndex(code.slice(mockResult.index!))
-        const [, declaration, path] = mockResult
-
-        if (lastIndex === null) continue
-
-        const startIndex = mockResult.index!
-
-        const { insideComment, insideString } = getIndexStatus(code, startIndex)
-
-        if (insideComment || insideString)
-          continue
-
-        const endIndex = startIndex + lastIndex
-
-        const filepath = await this.resolve(path, id)
-
-        m ??= new MagicString(code)
-
-        const overwrite = getMethodCall('mock', filepath?.id || path, path)
-
-        m.overwrite(startIndex, startIndex + declaration.length, overwrite)
-        m.prepend(`${m.slice(startIndex, endIndex)}\n`)
-        m.remove(startIndex, endIndex)
-      }
+    async transform(code) {
+      const m = hoistMocks(code)
 
       if (m) {
         // hoist vitest imports in case it was used inside vi.mock factory #425
@@ -80,14 +64,6 @@ function getMockLastIndex(code: string): number | null {
   return code[index + 1] === ';' ? index + 2 : index + 1
 }
 
-function getMethodCall(method: string, actualPath: string, importPath: string) {
-  let nodeModule = 'null'
-  if (actualPath.includes('/node_modules/'))
-    nodeModule = `"${importPath}"`
-
-  return `__vitest__${method}__("${actualPath}", ${nodeModule}`
-}
-
 function getIndexStatus(code: string, from: number) {
   let index = 0
   let commentStarted = false