From ca9507d9d0501b4f10bd5b18fd7ed90353e23fa9 Mon Sep 17 00:00:00 2001
From: Vladimir Sheremet <sheremet.va@icloud.com>
Date: Mon, 17 Jan 2022 14:47:56 +0300
Subject: [PATCH] refactor: mocker refactor

---
 packages/vitest/src/node/execute.ts | 32 ++++++++++-----------
 packages/vitest/src/node/mocker.ts  | 44 ++++++++++++++++++-----------
 2 files changed, 43 insertions(+), 33 deletions(-)

diff --git a/packages/vitest/src/node/execute.ts b/packages/vitest/src/node/execute.ts
index 0c0e69998..6fabc0e7b 100644
--- a/packages/vitest/src/node/execute.ts
+++ b/packages/vitest/src/node/execute.ts
@@ -42,32 +42,29 @@ export class VitestRunner extends ViteNodeRunner {
   }
 
   prepareContext(context: Record<string, any>) {
-    const suite = this.mocker.getSuiteFilepath()
-    const mockMap = this.options.mockMap
     const request = context.__vite_ssr_import__
 
     const callFunctionMock = async(dep: string, mock: () => any) => {
-      const name = `${dep}__mock`
-      const cached = this.moduleCache.get(name)?.exports
+      const cacheName = `${dep}__mock`
+      const cached = this.moduleCache.get(cacheName)?.exports
       if (cached)
         return cached
       const exports = await mock()
-      this.setCache(name, { exports })
+      this.setCache(cacheName, { exports })
       return exports
     }
 
     const requestWithMock = async(dep: string) => {
-      const mocks = mockMap[suite || ''] || {}
-      const mock = mocks[this.mocker.resolveDependency(dep)]
+      const mock = this.mocker.getDependencyMock(dep)
       if (mock === null) {
-        const mockedKey = `${dep}__mock`
-        const cache = this.moduleCache.get(mockedKey)
+        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(mockedKey, { exports })
+        this.setCache(cacheName, { exports })
         return exports
       }
       if (typeof mock === 'function')
@@ -76,16 +73,17 @@ export class VitestRunner extends ViteNodeRunner {
         dep = mock
       return request(dep)
     }
-    const importActual = (path: string, nmName: string) => {
-      return request(this.mocker.getActualPath(path, nmName))
+    const importActual = (path: string, external: string | null) => {
+      return request(this.mocker.getActualPath(path, external))
     }
-    const importMock = async(path: string, nmName: string): Promise<any> => {
-      if (!suite)
-        throw new Error('You can import mock only inside of a running test')
+    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)
 
-      const mock = (mockMap[suite] || {})[path] || this.mocker.resolveMockPath(path, this.root, nmName)
       if (mock === null) {
-        const fsPath = this.mocker.getActualPath(path, nmName)
+        const fsPath = this.mocker.getActualPath(path, external)
         const mod = await request(fsPath)
         return this.mocker.mockObject(mod)
       }
diff --git a/packages/vitest/src/node/mocker.ts b/packages/vitest/src/node/mocker.ts
index 67d7e56f7..a1e07d905 100644
--- a/packages/vitest/src/node/mocker.ts
+++ b/packages/vitest/src/node/mocker.ts
@@ -4,7 +4,7 @@ import { basename, dirname, join, resolve } from 'pathe'
 import { spies, spyOn } from '../integrations/jest-mock'
 import { mergeSlashes, normalizeId } from '../utils'
 
-export type SuiteMocks = Record<string, Record<string, string | null | (() => any)>>
+export type SuiteMocks = Record<string, Record<string, string | null | (() => unknown)>>
 
 function resolveMockPath(mockPath: string, root: string, external: string | null) {
   const path = normalizeId(external || mockPath)
@@ -85,34 +85,45 @@ function mockObject(obj: any) {
 
 export function createMocker(root: string, mockMap: SuiteMocks) {
   function getSuiteFilepath() {
-    return process.__vitest_worker__?.filepath
+    return process.__vitest_worker__?.filepath || 'global'
   }
 
-  function getActualPath(path: string, external: string) {
+  function getMocks() {
+    const suite = getSuiteFilepath()
+    const suiteMocks = mockMap[suite || '']
+    const globalMocks = mockMap.global
+
+    return {
+      ...suiteMocks,
+      ...globalMocks,
+    }
+  }
+
+  function getDependencyMock(dep: string) {
+    return getMocks()[resolveDependency(dep)]
+  }
+
+  function getActualPath(path: string, external: string | null) {
     if (external)
       return mergeSlashes(`/@fs/${path}`)
 
     return normalizeId(path.replace(root, ''))
   }
 
-  function unmockPath(path: string, nmName: string) {
+  function unmockPath(path: string, external: string | null) {
     const suitefile = getSuiteFilepath()
 
-    if (suitefile) {
-      const fsPath = getActualPath(path, nmName)
-      mockMap[suitefile] ??= {}
-      delete mockMap[suitefile][fsPath]
-    }
+    const fsPath = getActualPath(path, external)
+    mockMap[suitefile] ??= {}
+    delete mockMap[suitefile][fsPath]
   }
 
-  function mockPath(path: string, nmName: string, factory?: () => any) {
+  function mockPath(path: string, external: string | null, factory?: () => any) {
     const suitefile = getSuiteFilepath()
 
-    if (suitefile) {
-      const fsPath = getActualPath(path, nmName)
-      mockMap[suitefile] ??= {}
-      mockMap[suitefile][fsPath] = factory || resolveMockPath(path, root, nmName)
-    }
+    const fsPath = getActualPath(path, external)
+    mockMap[suitefile] ??= {}
+    mockMap[suitefile][fsPath] = factory || resolveMockPath(path, root, external)
   }
 
   function clearMocks({ clearMocks, mockReset, restoreMocks }: { clearMocks: boolean; mockReset: boolean; restoreMocks: boolean }) {
@@ -142,10 +153,11 @@ export function createMocker(root: string, mockMap: SuiteMocks) {
     unmockPath,
     clearMocks,
     getActualPath,
+    getMocks,
+    getDependencyMock,
 
     mockObject,
     getSuiteFilepath,
     resolveMockPath,
-    resolveDependency,
   }
 }
-- 
2.51.2