From 7719e79ee5b36abeb8bf55e108dc049d5af27293 Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Fri, 5 Jan 2024 09:08:24 +0100
Subject: [PATCH] fix(vitest): vi.mock breaks tests when using imported
 variables inside the factory (#4873)

Co-authored-by: Dunqing <dengqing0821@gmail.com>
---
 packages/vitest/src/node/hoistMocks.ts | 35 ++++++++-------
 test/core/test/injector-mock.test.ts   | 60 +++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 16 deletions(-)

diff --git a/packages/vitest/src/node/hoistMocks.ts b/packages/vitest/src/node/hoistMocks.ts
index d279196ba3c3..61caeecda49e 100644
--- a/packages/vitest/src/node/hoistMocks.ts
+++ b/packages/vitest/src/node/hoistMocks.ts
@@ -78,7 +78,6 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
 
   const hoistIndex = code.match(hashbangRE)?.[0].length ?? 0
 
-  let hoistedCode = ''
   let hoistedVitestImports = ''
 
   let uid = 0
@@ -148,6 +147,7 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
   }
 
   const declaredConst = new Set<string>()
+  const hoistedNodes: Node[] = []
 
   esmWalker(ast, {
     onIdentifier(id, info, parentStack) {
@@ -184,12 +184,8 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
       ) {
         const methodName = node.callee.property.name
 
-        if (methodName === 'mock' || methodName === 'unmock') {
-          const end = getBetterEnd(code, node)
-          const nodeCode = code.slice(node.start, end)
-          hoistedCode += `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
-          s.remove(node.start, end)
-        }
+        if (methodName === 'mock' || methodName === 'unmock')
+          hoistedNodes.push(node)
 
         if (methodName === 'hoisted') {
           const declarationNode = findNodeAround(ast, node.start, 'VariableDeclaration')?.node as Positioned<VariableDeclaration> | undefined
@@ -212,23 +208,32 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
 
           if (canMoveDeclaration) {
             // hoist "const variable = vi.hoisted(() => {})"
-            const end = getBetterEnd(code, declarationNode)
-            const nodeCode = code.slice(declarationNode.start, end)
-            hoistedCode += `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
-            s.remove(declarationNode.start, end)
+            hoistedNodes.push(declarationNode)
           }
           else {
             // hoist "vi.hoisted(() => {})"
-            const end = getBetterEnd(code, node)
-            const nodeCode = code.slice(node.start, end)
-            hoistedCode += `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
-            s.remove(node.start, end)
+            hoistedNodes.push(node)
           }
         }
       }
     },
   })
 
+  // Wait for imports to be hoisted and then hoist the mocks
+  const hoistedCode = hoistedNodes.map((node) => {
+    const end = getBetterEnd(code, node)
+    /**
+     * In the following case, we need to change the `user` to user: __vi_import_x__.user
+     * So we should get the latest code from `s`.
+     *
+     * import user from './user'
+     * vi.mock('./mock.js', () => ({ getSession: vi.fn().mockImplementation(() => ({ user })) }))
+     */
+    const nodeCode = s.slice(node.start, end)
+    s.remove(node.start, end)
+    return `${nodeCode}${nodeCode.endsWith('\n') ? '' : '\n'}`
+  }).join('')
+
   if (hoistedCode || hoistedVitestImports) {
     s.prepend(
       hoistedVitestImports
diff --git a/test/core/test/injector-mock.test.ts b/test/core/test/injector-mock.test.ts
index 705c6e9a2c09..28e058217ec1 100644
--- a/test/core/test/injector-mock.test.ts
+++ b/test/core/test/injector-mock.test.ts
@@ -99,7 +99,7 @@ describe('transform', () => {
   }
   test('default import', async () => {
     expect(
-      await hoistSimpleCodeWithoutMocks(`import foo from 'vue';console.log(foo.bar)`),
+      hoistSimpleCodeWithoutMocks(`import foo from 'vue';console.log(foo.bar)`),
     ).toMatchInlineSnapshot(`
       "const { vi } = await import('vitest')
       vi.mock('faker');
@@ -109,6 +109,64 @@ describe('transform', () => {
     `)
   })
 
+  test('can use imported variables inside the mock', () => {
+    expect(
+      hoistMocks(`
+import { vi } from 'vitest'
+import user from './user'
+import { admin } from './admin'
+vi.mock('./mock.js', () => ({
+  getSession: vi.fn().mockImplementation(() => ({
+    user,
+    admin: admin,
+  }))
+}))
+`, './test.js', parse)?.code.trim(),
+    ).toMatchInlineSnapshot(`
+      "const { vi } = await import('vitest')
+      vi.mock('./mock.js', () => ({
+        getSession: vi.fn().mockImplementation(() => ({
+          user: __vi_import_0__.default,
+          admin: __vi_import_1__.admin,
+        }))
+      }))
+      const __vi_import_0__ = await import('./user')
+      const __vi_import_1__ = await import('./admin')"
+    `)
+  })
+
+  test('can use hoisted variables inside the mock', () => {
+    expect(
+      hoistMocks(`
+import { vi } from 'vitest'
+const { user, admin } = await vi.hoisted(async () => {
+  const { default: user } = await import('./user')
+  const { admin } = await import('./admin')
+  return { user, admin }
+})
+vi.mock('./mock.js', () => {
+  getSession: vi.fn().mockImplementation(() => ({
+    user,
+    admin: admin,
+  }))
+})
+`, './test.js', parse)?.code.trim(),
+    ).toMatchInlineSnapshot(`
+      "const { vi } = await import('vitest')
+      const { user, admin } = await vi.hoisted(async () => {
+        const { default: user } = await import('./user')
+        const { admin } = await import('./admin')
+        return { user, admin }
+      })
+      vi.mock('./mock.js', () => {
+        getSession: vi.fn().mockImplementation(() => ({
+          user,
+          admin: admin,
+        }))
+      })"
+    `)
+  })
+
   test('named import', async () => {
     expect(
       await hoistSimpleCodeWithoutMocks(