From 65d71b9e97e6d4a6d7e38e12ed2c436e41dc97e3 Mon Sep 17 00:00:00 2001
From: Wenzhe Wang <mysteryven@gmail.com>
Date: Wed, 22 Feb 2023 19:52:43 +0800
Subject: [PATCH] fix: cleanup last mocked cache when call vi.doMock (#2872)

fix https://github.com/vitest-dev/vitest/issues/2870
---
 docs/api/vi.md                        |  6 ++---
 packages/vitest/src/runtime/mocker.ts | 11 ++++++---
 test/core/test/do-mock.test.ts        | 32 +++++++++++++++++++++++++++
 test/core/test/fixtures/increment.ts  |  1 +
 4 files changed, 43 insertions(+), 7 deletions(-)
 create mode 100644 test/core/test/do-mock.test.ts
 create mode 100644 test/core/test/fixtures/increment.ts

diff --git a/docs/api/vi.md b/docs/api/vi.md
index 8c48ed880640..dfc44e4f9cbe 100644
--- a/docs/api/vi.md
+++ b/docs/api/vi.md
@@ -221,13 +221,11 @@ increment(1) === 2
 let mockedIncrement = 100
 
 beforeEach(() => {
-  // simple doMock doesn't clear the previous cache, so we need to clear it manually here
-  vi.doUnmock('./increment.js')
   // you can access variables inside a factory
-  vi.doMock('./increment.js', () => ({ increment: () => mockedIncrement++ }))
+  vi.doMock('./increment.js', () => ({ increment: () => ++mockedIncrement }))
 })
 
-test('importing the next module imports mocked one', () => {
+test('importing the next module imports mocked one', async () => {
   // original import WAS NOT MOCKED, because vi.doMock is evaluated AFTER imports
   expect(increment(1)).toBe(2)
   const { increment: mockedIncrement } = await import('./increment.js')
diff --git a/packages/vitest/src/runtime/mocker.ts b/packages/vitest/src/runtime/mocker.ts
index 8502420908a4..4a770fc10bb6 100644
--- a/packages/vitest/src/runtime/mocker.ts
+++ b/packages/vitest/src/runtime/mocker.ts
@@ -60,6 +60,12 @@ export class VitestMocker {
     return this.executor.moduleCache
   }
 
+  private deleteCachedItem(id: string) {
+    const mockId = this.getMockPath(id)
+    if (this.moduleCache.has(mockId))
+      this.moduleCache.delete(mockId)
+  }
+
   public getSuiteFilepath(): string {
     return getWorkerState().filepath || 'global'
   }
@@ -292,9 +298,7 @@ export class VitestMocker {
     if (mock && id in mock)
       delete mock[id]
 
-    const mockId = this.getMockPath(id)
-    if (this.moduleCache.get(mockId))
-      this.moduleCache.delete(mockId)
+    this.deleteCachedItem(id)
   }
 
   public mockPath(originalId: string, path: string, external: string | null, factory?: MockFactory) {
@@ -309,6 +313,7 @@ export class VitestMocker {
 
     this.mockMap.set(suitefile, mocks)
     this.resolveCache.set(suitefile, resolves)
+    this.deleteCachedItem(id)
   }
 
   public async importActual<T>(rawId: string, importee: string): Promise<T> {
diff --git a/test/core/test/do-mock.test.ts b/test/core/test/do-mock.test.ts
new file mode 100644
index 000000000000..4cec6e845ac7
--- /dev/null
+++ b/test/core/test/do-mock.test.ts
@@ -0,0 +1,32 @@
+import { expect, test, vi } from 'vitest'
+
+test('doMock works', async () => {
+  const { increment: incrementWith1 } = await import('./fixtures/increment')
+  expect(incrementWith1(1)).toBe(2)
+
+  vi.doMock('./fixtures/increment', () => ({
+    increment: (num: number) => num + 10,
+  }))
+
+  const { increment: incrementWith10 } = await import('./fixtures/increment')
+
+  expect(incrementWith10(1)).toBe(11)
+})
+
+test('the second doMock can override the first doMock', async () => {
+  vi.doMock('./fixtures/increment', () => ({
+    increment: (num: number) => num + 10,
+  }))
+
+  const { increment: incrementWith1 } = await import('./fixtures/increment')
+
+  expect(incrementWith1(1)).toBe(11)
+
+  vi.doMock('./fixtures/increment', () => ({
+    increment: (num: number) => num + 20,
+  }))
+
+  const { increment: incrementWith20 } = await import('./fixtures/increment')
+
+  expect(incrementWith20(1)).toBe(21)
+})
diff --git a/test/core/test/fixtures/increment.ts b/test/core/test/fixtures/increment.ts
new file mode 100644
index 000000000000..305399b25cfe
--- /dev/null
+++ b/test/core/test/fixtures/increment.ts
@@ -0,0 +1 @@
+export const increment = (num: number) => num + 1