fix(vitest): don't hang when mocking files with cyclic dependencies (#4811)

Author: sheremet-va

examples/mocks/src/cyclic-deps/module-1.mjs

@@ -0,0 +1 @@
+import './module-2'

examples/mocks/src/cyclic-deps/module-2.mjs

@@ -0,0 +1 @@
+import './module-3'

examples/mocks/src/cyclic-deps/module-3.mjs

@@ -0,0 +1 @@
+import './module-4'

examples/mocks/src/cyclic-deps/module-4.mjs

@@ -0,0 +1 @@
+import './module-1'

examples/mocks/test/cyclic-import-actual.spec.ts

@@ -0,0 +1,13 @@
+import { expect, test, vi } from 'vitest'
+
+import '../src/cyclic-deps/module-1'
+
+vi.mock('../src/cyclic-deps/module-2', async () => {
+  await vi.importActual('../src/cyclic-deps/module-2')
+
+  return { default: () => {} }
+})
+
+test('some test', () => {
+  expect(1 + 1).toBe(2)
+})

examples/mocks/test/cyclic-import-original.spec.ts

@@ -0,0 +1,13 @@
+import { expect, test, vi } from 'vitest'
+
+import '../src/cyclic-deps/module-1'
+
+vi.mock('../src/cyclic-deps/module-2', async (importOriginal) => {
+  await importOriginal()
+
+  return { default: () => {} }
+})
+
+test('some test', () => {
+  expect(1 + 1).toBe(2)
+})

packages/vitest/src/integrations/vi.ts

@@ -482,7 +482,7 @@ function createVitest(): VitestUtils {
       _mocker.queueMock(
         path,
         importer,
-        factory ? () => factory(() => _mocker.importActual(path, importer)) : undefined,
+        factory ? () => factory(() => _mocker.importActual(path, importer, _mocker.getMockContext().callstack)) : undefined,
       )
     },
 
@@ -495,7 +495,7 @@ function createVitest(): VitestUtils {
       _mocker.queueMock(
         path,
         importer,
-        factory ? () => factory(() => _mocker.importActual(path, importer)) : undefined,
+        factory ? () => factory(() => _mocker.importActual(path, importer, _mocker.getMockContext().callstack)) : undefined,
       )
     },
 
@@ -504,7 +504,11 @@ function createVitest(): VitestUtils {
     },
 
     async importActual<T = unknown>(path: string): Promise<T> {
-      return _mocker.importActual<T>(path, getImporter())
+      return _mocker.importActual<T>(
+        path,
+        getImporter(),
+        _mocker.getMockContext().callstack,
+      )
     },
 
     async importMock<T>(path: string): Promise<MaybeMockedDeep<T>> {

packages/vitest/src/runtime/mocker.ts

@@ -32,6 +32,13 @@ class RefTracker {
 
 type Key = string | symbol
 
+interface MockContext {
+  /**
+   * When mocking with a factory, this refers to the module that imported the mock.
+   */
+  callstack: null | string[]
+}
+
 function isSpecialProp(prop: Key, parentType: string) {
   return parentType.includes('Function')
       && typeof prop === 'string'
@@ -54,6 +61,10 @@ export class VitestMocker {
 
   private filterPublicKeys: (symbol | string)[]
 
+  private mockContext: MockContext = {
+    callstack: null,
+  }
+
   constructor(
     public executor: VitestExecutor,
   ) {
@@ -201,9 +212,9 @@ export class VitestMocker {
           throw this.createError(
             `[vitest] No "${String(prop)}" export is defined on the "${mockpath}" mock. `
             + 'Did you forget to return it from "vi.mock"?'
-            + '\nIf you need to partially mock a module, you can use "vi.importActual" inside:\n\n'
-            + `${c.green(`vi.mock("${mockpath}", async () => {
-  const actual = await vi.importActual("${mockpath}")
+            + '\nIf you need to partially mock a module, you can use "importOriginal" helper inside:\n\n'
+            + `${c.green(`vi.mock("${mockpath}", async (importOriginal) => {
+  const actual = await importOriginal()
   return {
     ...actual,
     // your mocked methods
@@ -221,6 +232,10 @@ export class VitestMocker {
     return moduleExports
   }
 
+  public getMockContext() {
+    return this.mockContext
+  }
+
   public getMockPath(dep: string) {
     return `mock:${dep}`
   }
@@ -407,9 +422,9 @@ export class VitestMocker {
     this.deleteCachedItem(id)
   }
 
-  public async importActual<T>(rawId: string, importee: string): Promise<T> {
-    const { id, fsPath } = await this.resolvePath(rawId, importee)
-    const result = await this.executor.cachedRequest(id, fsPath, [importee])
+  public async importActual<T>(rawId: string, importer: string, callstack?: string[] | null): Promise<T> {
+    const { id, fsPath } = await this.resolvePath(rawId, importer)
+    const result = await this.executor.cachedRequest(id, fsPath, callstack || [importer])
     return result as T
   }
 
@@ -453,9 +468,14 @@ export class VitestMocker {
     if (typeof mock === 'function' && !callstack.includes(mockPath) && !callstack.includes(url)) {
       try {
         callstack.push(mockPath)
+        // this will not work if user does Promise.all(import(), import())
+        // we can also use AsyncLocalStorage to store callstack, but this won't work in the browser
+        // maybe we should improve mock API in the future?
+        this.mockContext.callstack = callstack
         return await this.callFunctionMock(mockPath, mock)
       }
       finally {
+        this.mockContext.callstack = null
         const indexMock = callstack.indexOf(mockPath)
         callstack.splice(indexMock, 1)
       }