fix(vitest-pool-workers): reset module graphs to ensure module mock works in watch mode (#8182)

Co-authored-by: Carmen Popoviciu <cpopoviciu@cloudflare.com>

Author: edmundhung

.changeset/quick-shoes-jump.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/vitest-pool-workers": patch
+---
+
+fix mock modules not re-evaluated in watch mode

packages/vitest-pool-workers/src/worker/lib/cloudflare/test-runner.ts

@@ -8,6 +8,7 @@ import {
 	registerHandlerAndGlobalWaitUntil,
 	waitForGlobalWaitUntil,
 } from "cloudflare:test-internal";
+import { vi } from "vitest";
 import { VitestTestRunner } from "vitest/runners";
 import workerdUnsafe from "workerd:unsafe";
 import type { CancelReason, Suite, Test } from "@vitest/runner";
@@ -229,6 +230,11 @@ export default class WorkersTestRunner extends VitestTestRunner {
 			await scheduler.wait(100);
 		}
 
+		// Unlike the official threads and forks pool, we do not recycle the miniflare instances to maintain the module cache.
+		// However, this creates a side effect where the module mock will not be re-evaluated on watch mode.
+		// This fixes https://github.com/cloudflare/workers-sdk/issues/6844 by resetting the module graph.
+		vi.resetModules();
+
 		// Ensure all `ctx.waitUntil()` calls complete before disposing the runtime
 		// (if using `vitest run`) and aborting all objects. `ctx.waitUntil()`s may
 		// contain storage calls (e.g. caching responses) that could try to access

packages/vitest-pool-workers/test/watch.test.ts

@@ -106,3 +106,103 @@ test("automatically re-runs integration tests", async ({
 		expect(result.stdout).toMatch("Tests  1 passed");
 	});
 });
+
+test("automatically reset module graph", async ({
+	expect,
+	seed,
+	vitestDev,
+}) => {
+	await seed({
+		"vitest.config.mts": dedent`
+			import { defineWorkersConfig } from "@cloudflare/vitest-pool-workers/config";
+			export default defineWorkersConfig({
+				test: {
+					poolOptions: {
+						workers: {
+							main: "./index.ts",
+							singleWorker: true,
+							miniflare: {
+								compatibilityDate: "2024-01-01",
+								compatibilityFlags: ["nodejs_compat"],
+							},
+						},
+					},
+				}
+			});
+		`,
+		"answer.ts": dedent`
+			export function getAnswer() {
+				return "wrong";
+			}
+		`,
+		"index.ts": dedent`
+			import { getAnswer } from "./answer";
+
+			export default {
+				async fetch(request, env, ctx) {
+					const answer = getAnswer();
+					return new Response(answer);
+				}
+			}
+		`,
+		"index.test.ts": dedent`
+			import { env, createExecutionContext, waitOnExecutionContext } from "cloudflare:test";
+			import { it, expect, vi } from "vitest";
+			import worker from "./index";
+			import { getAnswer } from './answer';
+
+			vi.mock('./answer');
+
+			it("mocks module properly", async () => {
+				vi.mocked(getAnswer).mockReturnValue("correct");
+
+				const request = new Request("https://example.com");
+				const ctx = createExecutionContext();
+				const response = await worker.fetch(request, env, ctx);
+				await waitOnExecutionContext(ctx);
+				expect(await response.text()).toBe("correct");
+			});
+		`,
+	});
+	const result = vitestDev();
+
+	await waitFor(() => {
+		expect(result.stdout).toMatch("Tests  1 passed");
+	});
+
+	// Trigger a re-run by updating the test file with an extra test.
+	await seed({
+		"index.test.ts": dedent`
+			import { env, createExecutionContext, waitOnExecutionContext } from "cloudflare:test";
+			import { it, expect, vi } from "vitest";
+			import worker from "./index";
+			import { getAnswer } from './answer';
+
+			vi.mock('./answer');
+
+			it("mocks module properly", async () => {
+				vi.mocked(getAnswer).mockReturnValue("correct");
+
+				const request = new Request("https://example.com");
+				const ctx = createExecutionContext();
+				const response = await worker.fetch(request, env, ctx);
+				await waitOnExecutionContext(ctx);
+				expect(await response.text()).toBe("correct");
+			});
+
+			it("mocks module properly when re-run in watch mode", async () => {
+				vi.mocked(getAnswer).mockReturnValue("test");
+
+				const request = new Request("https://example.com");
+				const ctx = createExecutionContext();
+				const response = await worker.fetch(request, env, ctx);
+				await waitOnExecutionContext(ctx);
+				expect(await response.text()).toBe("test");
+			});
+		`,
+	});
+
+	await waitFor(() => {
+		expect(result.stdout).toMatch("Tests  2 passed");
+	});
+});