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,17 +1,25 @@
 /* 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'
 
 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,
     })
   }
 }