From 78cd34b17bd7647d898edf23ba85a16c95500175 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sat, 30 May 2026 13:44:54 +0900 Subject: [PATCH] fix(testing): isolate bun module mocks across tests --- .../abort-with-timeout.test.ts | 6 +- src/testing/module-mock-lifecycle-bun.test.ts | 161 +++++++++++++ src/testing/module-mock-lifecycle-consumer.ts | 1 + src/testing/module-mock-lifecycle-fixture.ts | 1 + .../module-mock-lifecycle-nested/caller.ts | 7 + .../module-mock-lifecycle-nested/fixture.ts | 1 + .../module-mock-lifecycle-tracking.test.ts | 155 +++++++++++++ src/testing/module-mock-lifecycle.test.ts | 90 +++++++- src/testing/module-mock-lifecycle.ts | 216 +++++++++++++----- src/testing/module-mock-stack.ts | 61 +++++ test-setup.ts | 6 +- 11 files changed, 642 insertions(+), 63 deletions(-) create mode 100644 src/testing/module-mock-lifecycle-bun.test.ts create mode 100644 src/testing/module-mock-lifecycle-consumer.ts create mode 100644 src/testing/module-mock-lifecycle-fixture.ts create mode 100644 src/testing/module-mock-lifecycle-nested/caller.ts create mode 100644 src/testing/module-mock-lifecycle-nested/fixture.ts create mode 100644 src/testing/module-mock-lifecycle-tracking.test.ts create mode 100644 src/testing/module-mock-stack.ts diff --git a/src/features/background-agent/abort-with-timeout.test.ts b/src/features/background-agent/abort-with-timeout.test.ts index 46b34e917..c0f3787af 100644 --- a/src/features/background-agent/abort-with-timeout.test.ts +++ b/src/features/background-agent/abort-with-timeout.test.ts @@ -1,4 +1,4 @@ -import { afterAll, describe, expect, mock, test } from "bun:test" +import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test" const logMock = mock(() => {}) @@ -20,6 +20,10 @@ function createClient(abort: (...args: Array) => Promise): Ope } describe("abortWithTimeout", () => { + beforeEach(() => { + logMock.mockClear() + }) + afterAll(() => { mock.restore() }) diff --git a/src/testing/module-mock-lifecycle-bun.test.ts b/src/testing/module-mock-lifecycle-bun.test.ts new file mode 100644 index 000000000..fdcdb6c16 --- /dev/null +++ b/src/testing/module-mock-lifecycle-bun.test.ts @@ -0,0 +1,161 @@ +/// + +import { describe, expect, mock, test } from "bun:test" +import { installModuleMockLifecycle } from "./module-mock-lifecycle" +import { mockNestedFixture } from "./module-mock-lifecycle-nested/caller" + +describe("installModuleMockLifecycle Bun integration", () => { + test("restores in-test ESM mock bindings during cleanup", async () => { + // given + mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked-aftereach" })) + const mockedModule = await import("./module-mock-lifecycle-fixture") + + // when + mock.restore() + + // then + expect(mockedModule.named).toBe("original") + }) + + test("keeps restored module mocks restored after test setup cleanup", async () => { + // when + const restoredModule = await import(`./module-mock-lifecycle-fixture?aftereach=${Date.now()}`) + + // then + expect(restoredModule.named).toBe("original") + }) + + test("uses Bun resolution so extensionless TypeScript imports capture the loaded module", () => { + // given + const moduleCalls: Array<{ specifier: string; value: Record }> = [] + const mockApi = { + module: (specifier: string, factory: () => Record) => { + moduleCalls.push({ specifier, value: factory() }) + }, + restore: mock(() => {}), + } + + installModuleMockLifecycle(mockApi) + + // when + mockApi.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" })) + mockApi.restore() + + // then + const restoreCall = moduleCalls.at(1) + expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-fixture\.ts$/) + expect(restoreCall?.value.named).toBe("original") + }) + + test("resolves relative mocks from the helper caller instead of the lifecycle module", () => { + // given + const moduleCalls: Array<{ specifier: string; value: Record }> = [] + const mockApi = { + module: (specifier: string, factory: () => Record) => { + moduleCalls.push({ specifier, value: factory() }) + }, + restore: mock(() => {}), + } + + installModuleMockLifecycle(mockApi) + + // when + mockNestedFixture(mockApi) + mockApi.restore() + + // then + const restoreCall = moduleCalls.at(1) + expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-nested\/fixture\.ts$/) + expect(restoreCall?.value.named).toBe("original") + }) + + test("restores Bun module mocks for later cache-busted imports", async () => { + // given + mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" })) + const mockedModule = await import("./module-mock-lifecycle-fixture") + + // when + mock.restore() + const restoredModule = await import(`./module-mock-lifecycle-fixture?restored=${Date.now()}`) + + // then + expect(mockedModule.named).toBe("original") + expect(restoredModule.named).toBe("original") + }) + + test("keeps Bun module mocks restored when cleanup calls mock.restore twice", async () => { + // given + mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" })) + const mockedModule = await import("./module-mock-lifecycle-fixture") + + // when + mock.restore() + mock.restore() + const restoredModule = await import(`./module-mock-lifecycle-fixture?double-restore=${Date.now()}`) + + // then + expect(mockedModule.named).toBe("original") + expect(restoredModule.named).toBe("original") + }) + + test("allows a later relative mock after restoring a resolved original", async () => { + // given + mock.module("./module-mock-lifecycle-fixture", () => ({ named: "first mock" })) + const firstMockedModule = await import("./module-mock-lifecycle-fixture") + + // when + mock.restore() + mock.module("./module-mock-lifecycle-fixture", () => ({ named: "second mock" })) + const secondMockedModule = await import("./module-mock-lifecycle-fixture") + + // then + expect(firstMockedModule.named).toBe("second mock") + expect(secondMockedModule.named).toBe("second mock") + }) + + test("allows a later caller-relative mock after restoring a file-url original", async () => { + // given + mock.module("./module-mock-lifecycle-nested/fixture", () => ({ named: "first mock" })) + const firstMockedModule = await import("./module-mock-lifecycle-nested/fixture") + + // when + mock.restore() + mockNestedFixture(mock) + const secondMockedModule = await import("./module-mock-lifecycle-nested/fixture") + + // then + expect(firstMockedModule.named).toBe("mocked") + expect(secondMockedModule.named).toBe("mocked") + }) + + test("restores Bun module mocks outside the lifecycle helper directory", async () => { + // given + mock.module("../cli/run/json-output", () => ({ + createJsonOutputManager: () => ({ mocked: true }), + })) + const mockedModule = await import("../cli/run/json-output") + + // when + mock.restore() + const restoredModule = await import(`../cli/run/json-output?restored=${Date.now()}`) + + // then + expect("mocked" in mockedModule.createJsonOutputManager()).toBe(false) + expect("mocked" in restoredModule.createJsonOutputManager()).toBe(false) + expect(typeof restoredModule.createJsonOutputManager().redirectToStderr).toBe("function") + }) + + test("restores dependency mocks after a cache-busted consumer import", async () => { + // given + mock.module("./module-mock-lifecycle-fixture", () => ({ named: "mocked" })) + const consumerModule = await import(`./module-mock-lifecycle-consumer?mocked=${Date.now()}`) + + // when + mock.restore() + const restoredModule = await import(`./module-mock-lifecycle-fixture?restored-dep=${Date.now()}`) + + // then + expect(consumerModule.consumed).toBe("original") + expect(restoredModule.named).toBe("original") + }) +}) diff --git a/src/testing/module-mock-lifecycle-consumer.ts b/src/testing/module-mock-lifecycle-consumer.ts new file mode 100644 index 000000000..881c0ed52 --- /dev/null +++ b/src/testing/module-mock-lifecycle-consumer.ts @@ -0,0 +1 @@ +export { named as consumed } from "./module-mock-lifecycle-fixture" diff --git a/src/testing/module-mock-lifecycle-fixture.ts b/src/testing/module-mock-lifecycle-fixture.ts new file mode 100644 index 000000000..11cea8714 --- /dev/null +++ b/src/testing/module-mock-lifecycle-fixture.ts @@ -0,0 +1 @@ +export const named = "original" diff --git a/src/testing/module-mock-lifecycle-nested/caller.ts b/src/testing/module-mock-lifecycle-nested/caller.ts new file mode 100644 index 000000000..d4cfd4382 --- /dev/null +++ b/src/testing/module-mock-lifecycle-nested/caller.ts @@ -0,0 +1,7 @@ +type MockModuleApi = { + module: (specifier: string, factory: () => Record) => unknown +} + +export function mockNestedFixture(mockApi: MockModuleApi): void { + mockApi.module("./fixture", () => ({ named: "mocked" })) +} diff --git a/src/testing/module-mock-lifecycle-nested/fixture.ts b/src/testing/module-mock-lifecycle-nested/fixture.ts new file mode 100644 index 000000000..11cea8714 --- /dev/null +++ b/src/testing/module-mock-lifecycle-nested/fixture.ts @@ -0,0 +1 @@ +export const named = "original" diff --git a/src/testing/module-mock-lifecycle-tracking.test.ts b/src/testing/module-mock-lifecycle-tracking.test.ts new file mode 100644 index 000000000..fd8ba19b5 --- /dev/null +++ b/src/testing/module-mock-lifecycle-tracking.test.ts @@ -0,0 +1,155 @@ +/// + +import { describe, expect, mock, test } from "bun:test" +import { installModuleMockLifecycle } from "./module-mock-lifecycle" + +describe("installModuleMockLifecycle active-test tracking", () => { + test("reapplies module-evaluation mocks with their mock factory after active restore", () => { + // given + const events: string[] = [] + const loadOriginalModule = mock(() => ({ ok: true as const, value: { named: "original" } })) + const mockApi = { + module: (specifier: string, factory: () => Record) => { + events.push(`module:${specifier}:${String(factory().named)}`) + }, + restore: mock(() => { + events.push("delegate:restore") + }), + } + + const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, { + getCallerStack: () => "Error\n at file:///repo/tests/top-level.test.ts:5:1\n at moduleEvaluation (native:1:11)", + getCallerUrl: () => "file:///repo/tests/top-level.test.ts", + trackOnlyDuringActiveTest: true, + resolveSpecifier: (specifier) => `resolved:${specifier}`, + loadOriginalModule, + }) + + // when + mockApi.module("./dependency", () => ({ named: "mocked" })) + beginTestMockTracking() + mockApi.restore() + endTestMockTracking() + + // then + expect(events).toEqual([ + "module:./dependency:mocked", + "delegate:restore", + "module:./dependency:mocked", + "module:resolved:./dependency:mocked", + ]) + expect(loadOriginalModule).not.toHaveBeenCalled() + }) + + test("restores original exports for active in-test mocks", () => { + // given + const events: string[] = [] + const mockApi = { + module: (specifier: string, factory: () => Record) => { + events.push(`module:${specifier}:${String(factory().named)}`) + }, + restore: mock(() => { + events.push("delegate:restore") + }), + } + + const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, { + getCallerStack: () => "Error\n at file:///repo/tests/example.test.ts:5:1\n at test (native:1:11)", + getCallerUrl: () => "file:///repo/tests/example.test.ts", + trackOnlyDuringActiveTest: true, + resolveSpecifier: (specifier) => `resolved:${specifier}`, + loadOriginalModule: () => ({ ok: true, value: { named: "original" } }), + }) + + // when + beginTestMockTracking() + mockApi.module("./dependency", () => ({ named: "mocked" })) + mockApi.restore() + endTestMockTracking() + + // then + expect(events).toEqual([ + "module:./dependency:mocked", + "delegate:restore", + "module:./dependency:original", + "module:resolved:./dependency:original", + ]) + }) + + test("clears persistent module-evaluation snapshots when restore runs while inactive", () => { + // given + const events: string[] = [] + const mockApi = { + module: (specifier: string, factory: () => Record) => { + events.push(`module:${specifier}:${String(factory().named)}`) + }, + restore: mock(() => { + events.push("delegate:restore") + }), + } + + const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, { + getCallerStack: () => "Error\n at file:///repo/tests/top-level.test.ts:5:1\n at moduleEvaluation (native:1:11)", + getCallerUrl: () => "file:///repo/tests/top-level.test.ts", + trackOnlyDuringActiveTest: true, + resolveSpecifier: (specifier) => `resolved:${specifier}`, + loadOriginalModule: () => ({ ok: true, value: { named: "original" } }), + }) + + // when + mockApi.module("./dependency", () => ({ named: "mocked" })) + mockApi.restore() + beginTestMockTracking() + mockApi.restore() + endTestMockTracking() + + // then + expect(events).toEqual([ + "module:./dependency:mocked", + "delegate:restore", + "delegate:restore", + ]) + }) + + test("keeps unrelated module-evaluation mocks after inactive restore", () => { + // given + const events: string[] = [] + let callerStack = "Error\n at file:///repo/tests/first.test.ts:5:1\n at moduleEvaluation (native:1:11)" + let callerUrl = "file:///repo/tests/first.test.ts" + const mockApi = { + module: (specifier: string, factory: () => Record) => { + events.push(`module:${specifier}:${String(factory().named)}`) + }, + restore: mock(() => { + events.push("delegate:restore") + }), + } + + installModuleMockLifecycle(mockApi, { + getCallerStack: () => callerStack, + getCallerUrl: () => callerUrl, + trackOnlyDuringActiveTest: true, + resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`, + loadOriginalModule: () => ({ ok: true, value: { named: "original" } }), + }) + + mockApi.module("./first", () => ({ named: "first top-level" })) + callerStack = "Error\n at file:///repo/tests/second.test.ts:5:1\n at moduleEvaluation (native:1:11)" + callerUrl = "file:///repo/tests/second.test.ts" + mockApi.module("./second", () => ({ named: "second top-level" })) + + // when + callerStack = "Error\n at file:///repo/tests/first.test.ts:10:1\n at cleanup (native:1:11)" + callerUrl = "file:///repo/tests/first.test.ts" + mockApi.restore() + + // then + expect(events).toEqual([ + "module:./first:first top-level", + "module:./second:second top-level", + "delegate:restore", + "module:./second:second top-level", + "module:resolved:file:///repo/tests/second.test.ts:./second:second top-level", + ]) + }) +}) diff --git a/src/testing/module-mock-lifecycle.test.ts b/src/testing/module-mock-lifecycle.test.ts index f93029724..4b2bb9602 100644 --- a/src/testing/module-mock-lifecycle.test.ts +++ b/src/testing/module-mock-lifecycle.test.ts @@ -4,9 +4,12 @@ import { describe, expect, mock, test } from "bun:test" import { installModuleMockLifecycle } from "./module-mock-lifecycle" describe("installModuleMockLifecycle", () => { - test("restores the original module export object instead of a cloned snapshot", () => { + test("preserves plain export identity while cloning ESM namespace snapshots", () => { // given const originalExports = { named: "original" } + const moduleNamespaceExports = Object.defineProperty({ named: "namespace" }, Symbol.toStringTag, { + value: "Module", + }) const moduleCalls: Array<{ specifier: string; value: Record }> = [] const mockApi = { module: (specifier: string, factory: () => Record) => { @@ -18,17 +21,25 @@ describe("installModuleMockLifecycle", () => { installModuleMockLifecycle(mockApi, { getCallerUrl: () => "file:///repo/tests/example.test.ts", resolveSpecifier: (specifier) => `resolved:${specifier}`, - loadOriginalModule: () => ({ ok: true, value: originalExports }), + loadOriginalModule: (specifier) => ({ + ok: true, + value: specifier === "./namespace" ? moduleNamespaceExports : originalExports, + }), }) // when mockApi.module("./dependency", () => ({ named: "mocked" })) + mockApi.module("./namespace", () => ({ named: "mocked namespace" })) mockApi.restore() // then - expect(moduleCalls.map((call) => call.specifier)).toEqual(["./dependency", "resolved:./dependency"]) - const restoreCall = moduleCalls.find((call) => call.specifier === "resolved:./dependency") - expect(restoreCall?.value).toBe(originalExports) + const plainRestoreCall = moduleCalls.find((call) => call.specifier === "./dependency" && call.value.named === "original") + const namespaceRestoreCall = moduleCalls.find( + (call) => call.specifier === "resolved:./namespace" && call.value.named === "namespace", + ) + expect(plainRestoreCall?.value).toBe(originalExports) + expect(namespaceRestoreCall?.value).toEqual({ named: "namespace" }) + expect(namespaceRestoreCall?.value).not.toBe(moduleNamespaceExports) }) test("clears tracked snapshots after the delegate restore runs", () => { @@ -58,6 +69,42 @@ describe("installModuleMockLifecycle", () => { expect(events).toEqual([ "module:./dependency:mocked", "delegate:restore", + "module:./dependency:original", + "module:resolved:./dependency:original", + ]) + }) + + test("reapplies the last restore snapshot when mock.restore is called twice", () => { + // given + const events: string[] = [] + const mockApi = { + module: (specifier: string, factory: () => Record) => { + events.push(`module:${specifier}:${String(factory().named)}`) + }, + restore: mock(() => { + events.push("delegate:restore") + }), + } + + installModuleMockLifecycle(mockApi, { + getCallerUrl: () => "file:///repo/tests/example.test.ts", + resolveSpecifier: (specifier) => `resolved:${specifier}`, + loadOriginalModule: () => ({ ok: true, value: { named: "original" } }), + }) + + // when + mockApi.module("./dependency", () => ({ named: "mocked" })) + mockApi.restore() + mockApi.restore() + + // then + expect(events).toEqual([ + "module:./dependency:mocked", + "delegate:restore", + "module:./dependency:original", + "module:resolved:./dependency:original", + "delegate:restore", + "module:./dependency:original", "module:resolved:./dependency:original", ]) }) @@ -87,6 +134,37 @@ describe("installModuleMockLifecycle", () => { expect(loadCount).toBe(1) }) + test("restores each original mock specifier while sharing one resolved snapshot", () => { + // given + const moduleCalls: Array<{ specifier: string; value: Record }> = [] + const mockApi = { + module: (specifier: string, factory: () => Record) => { + moduleCalls.push({ specifier, value: factory() }) + }, + restore: mock(() => {}), + } + + installModuleMockLifecycle(mockApi, { + getCallerUrl: () => "file:///repo/tests/example.test.ts", + resolveSpecifier: () => "file:///repo/src/dependency.ts", + loadOriginalModule: () => ({ ok: true, value: { named: "original" } }), + }) + + // when + mockApi.module("../src/dependency", () => ({ named: "first mock" })) + mockApi.module("./dependency", () => ({ named: "second mock" })) + mockApi.restore() + + // then + expect(moduleCalls).toEqual([ + { specifier: "../src/dependency", value: { named: "first mock" } }, + { specifier: "./dependency", value: { named: "second mock" } }, + { specifier: "../src/dependency", value: { named: "original" } }, + { specifier: "file:///repo/src/dependency.ts", value: { named: "original" } }, + { specifier: "./dependency", value: { named: "original" } }, + ]) + }) + test("does not restore unresolved modules to avoid cleanup errors", () => { // given const moduleCalls: Array<{ specifier: string; value: Record }> = [] @@ -107,7 +185,7 @@ describe("installModuleMockLifecycle", () => { mockApi.module("virtual:missing", () => ({ named: "mocked" })) mockApi.restore() - // then - only the original mock call, no restore call for unresolved module + // then expect(moduleCalls).toEqual([{ specifier: "virtual:missing", value: { named: "mocked" } }]) }) }) diff --git a/src/testing/module-mock-lifecycle.ts b/src/testing/module-mock-lifecycle.ts index fce6203ea..c49cb25c0 100644 --- a/src/testing/module-mock-lifecycle.ts +++ b/src/testing/module-mock-lifecycle.ts @@ -1,5 +1,6 @@ import { createRequire } from "node:module" -import { pathToFileURL } from "node:url" +import { fileURLToPath } from "node:url" +import { defaultGetCallerStack, isModuleEvaluationStack, resolveCallerUrlFromStack } from "./module-mock-stack" type MockModuleFactory = () => Record @@ -13,28 +14,43 @@ type ModuleLoadResult = | { ok: false; error: Error } type ModuleSnapshot = { - restoreSpecifier: string + restoreOriginalSpecifiers: boolean + restoreSpecifiers: Set + restoreFactory: MockModuleFactory +} + +type PersistentModuleSnapshot = { + ownerUrls: Set + restoreSpecifiers: Set restoreFactory: MockModuleFactory } type ModuleMockLifecycleOptions = { + getCallerStack?: () => string getCallerUrl?: () => string + trackOnlyDuringActiveTest?: boolean resolveSpecifier?: (specifier: string, callerUrl: string) => string loadOriginalModule?: (specifier: string, callerUrl: string) => ModuleLoadResult } -function toError(error: unknown): Error { - if (error instanceof Error) { - return error - } +let originalLoadNonce = 0 - return new Error(String(error)) +function toError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)) } function isModuleExports(moduleValue: unknown): moduleValue is Record { return moduleValue !== null && typeof moduleValue === "object" } +function isModuleNamespaceObject(moduleValue: Record): boolean { + return Object.prototype.toString.call(moduleValue) === "[object Module]" +} + +function shouldRestoreOriginalSpecifier(moduleValue: unknown): boolean { + return !isModuleExports(moduleValue) || !isModuleNamespaceObject(moduleValue) +} + function createRestoreExports(moduleValue: unknown): Record { if (typeof moduleValue === "function") { const functionExports = Object.assign({}, moduleValue) @@ -45,47 +61,42 @@ function createRestoreExports(moduleValue: unknown): Record { } if (isModuleExports(moduleValue)) { + if (isModuleNamespaceObject(moduleValue)) { + return { ...moduleValue } + } + return moduleValue } return { default: moduleValue } } -function normalizeStackPath(rawPath: string): string { - if (rawPath.startsWith("file://")) { - return rawPath - } - - return pathToFileURL(rawPath).href +function resolveWithBun(specifier: string, callerUrl: string): string { + const callerDirectory = fileURLToPath(new URL(".", callerUrl)) + return Bun.resolveSync(specifier, callerDirectory) } -function defaultGetCallerUrl(): string { - const stack = new Error().stack ?? "" - const lines = stack.split("\n") - - for (const line of lines) { - const match = line.match(/(?:\()?(file:\/\/[^\s)]+|\/[^\s):]+):(\d+):(\d+)/) - const candidatePath = match?.[1] - if (!candidatePath) { - continue - } - - if ( - candidatePath.includes("/test-setup.ts") || - candidatePath.includes("/src/testing/module-mock-lifecycle.ts") - ) { - continue - } - - return normalizeStackPath(candidatePath) - } - - return import.meta.url +function isSchemeSpecifier(specifier: string): boolean { + return /^[a-zA-Z][a-zA-Z\d+.-]*:/.test(specifier) } function defaultResolveSpecifier(specifier: string, callerUrl: string): string { try { - return import.meta.resolve(specifier, callerUrl) + return resolveWithBun(specifier, callerUrl) + } catch { + return specifier + } +} + +function createOriginalLoadSpecifier(specifier: string, callerUrl: string): string { + try { + const resolved = resolveWithBun(specifier, callerUrl) + if (isSchemeSpecifier(resolved)) { + return specifier + } + + originalLoadNonce += 1 + return `${resolved}?omo_original=${originalLoadNonce}` } catch { return specifier } @@ -94,7 +105,7 @@ function defaultResolveSpecifier(specifier: string, callerUrl: string): string { function defaultLoadOriginalModule(specifier: string, callerUrl: string): ModuleLoadResult { try { const require = createRequire(callerUrl) - return { ok: true, value: require(specifier) } + return { ok: true, value: require(createOriginalLoadSpecifier(specifier, callerUrl)) } } catch (error) { return { ok: false, error: toError(error) } } @@ -103,46 +114,141 @@ function defaultLoadOriginalModule(specifier: string, callerUrl: string): Module export function installModuleMockLifecycle( mockApi: MockApi, options: ModuleMockLifecycleOptions = {}, -): { restoreModuleMocks: () => void } { +): { + beginTestMockTracking: () => void + endTestMockTracking: () => void + restoreModuleMocks: () => void +} { const snapshots = new Map() + const persistentSnapshots = new Map() + let lastRestoredSnapshots: ModuleSnapshot[] = [] + let isActiveTest = !options.trackOnlyDuringActiveTest const delegateModule = mockApi.module.bind(mockApi) const delegateRestore = mockApi.restore.bind(mockApi) - const getCallerUrl = options.getCallerUrl ?? defaultGetCallerUrl + const getCallerStack = options.getCallerStack ?? defaultGetCallerStack const resolveSpecifier = options.resolveSpecifier ?? defaultResolveSpecifier const loadOriginalModule = options.loadOriginalModule ?? defaultLoadOriginalModule - function restoreModuleMocks(): void { - for (const snapshot of snapshots.values()) { - delegateModule(snapshot.restoreSpecifier, snapshot.restoreFactory) + function getCallerUrl(callerStack: string): string { + return options.getCallerUrl?.() ?? resolveCallerUrlFromStack(callerStack) + } + + function restoreModuleMocksForRestoreCall(): void { + const snapshotsToRestore = snapshots.size > 0 ? Array.from(snapshots.values()) : lastRestoredSnapshots + + for (const snapshot of snapshotsToRestore) { + for (const restoreSpecifier of snapshot.restoreSpecifiers) { + delegateModule(restoreSpecifier, snapshot.restoreFactory) + } } - snapshots.clear() + if (snapshots.size > 0) { + lastRestoredSnapshots = snapshotsToRestore + snapshots.clear() + } + } + + function restorePersistentModuleMocksForRestoreCall(): void { + for (const snapshot of persistentSnapshots.values()) { + for (const restoreSpecifier of snapshot.restoreSpecifiers) { + delegateModule(restoreSpecifier, snapshot.restoreFactory) + } + } + } + + function clearPersistentModuleMocksForOwner(ownerUrl: string): void { + for (const [resolvedSpecifier, snapshot] of persistentSnapshots) { + snapshot.ownerUrls.delete(ownerUrl) + if (snapshot.ownerUrls.size === 0) { + persistentSnapshots.delete(resolvedSpecifier) + } + } + } + + function restoreModuleMocks(): void { + if (snapshots.size === 0) { + return + } + + restoreModuleMocksForRestoreCall() + } + + function beginTestMockTracking(): void { + isActiveTest = true + } + + function endTestMockTracking(): void { + isActiveTest = !options.trackOnlyDuringActiveTest } mockApi.module = (specifier: string, factory: MockModuleFactory): unknown => { - const callerUrl = getCallerUrl() - const restoreSpecifier = resolveSpecifier(specifier, callerUrl) + lastRestoredSnapshots = [] + const callerStack = getCallerStack() + const callerUrl = getCallerUrl(callerStack) + const isModuleEvaluation = isModuleEvaluationStack(callerStack) - if (!snapshots.has(restoreSpecifier)) { - const originalModule = loadOriginalModule(specifier, callerUrl) - - if (originalModule.ok) { - const restoreExports = createRestoreExports(originalModule.value) - snapshots.set(restoreSpecifier, { - restoreSpecifier, - restoreFactory: () => restoreExports, + if (isModuleEvaluation) { + const resolvedSpecifier = resolveSpecifier(specifier, callerUrl) + const existingSnapshot = persistentSnapshots.get(resolvedSpecifier) + if (existingSnapshot) { + existingSnapshot.ownerUrls.add(callerUrl) + existingSnapshot.restoreSpecifiers.add(specifier) + existingSnapshot.restoreSpecifiers.add(resolvedSpecifier) + } else { + persistentSnapshots.set(resolvedSpecifier, { + ownerUrls: new Set([callerUrl]), + restoreSpecifiers: new Set([specifier, resolvedSpecifier]), + restoreFactory: factory, }) } + return delegateModule(specifier, factory) + } + + if (isActiveTest) { + const resolvedSpecifier = resolveSpecifier(specifier, callerUrl) + const existingSnapshot = snapshots.get(resolvedSpecifier) + + if (existingSnapshot) { + if (existingSnapshot.restoreOriginalSpecifiers) { + existingSnapshot.restoreSpecifiers.add(specifier) + } + existingSnapshot.restoreSpecifiers.add(resolvedSpecifier) + } else { + const originalModule = loadOriginalModule(specifier, callerUrl) + + if (originalModule.ok) { + const restoreExports = createRestoreExports(originalModule.value) + const restoreOriginalSpecifiers = shouldRestoreOriginalSpecifier(originalModule.value) + snapshots.set(resolvedSpecifier, { + restoreOriginalSpecifiers, + restoreSpecifiers: new Set( + restoreOriginalSpecifiers ? [specifier, resolvedSpecifier] : [resolvedSpecifier], + ), + restoreFactory: () => restoreExports, + }) + } + } } return delegateModule(specifier, factory) } mockApi.restore = (): unknown => { + const callerStack = getCallerStack() + const callerUrl = getCallerUrl(callerStack) const result = delegateRestore() - restoreModuleMocks() + if (!isActiveTest) { + snapshots.clear() + lastRestoredSnapshots = [] + clearPersistentModuleMocksForOwner(callerUrl) + restorePersistentModuleMocksForRestoreCall() + return result + } + + restoreModuleMocksForRestoreCall() + restorePersistentModuleMocksForRestoreCall() return result } - return { restoreModuleMocks } + return { beginTestMockTracking, endTestMockTracking, restoreModuleMocks } } diff --git a/src/testing/module-mock-stack.ts b/src/testing/module-mock-stack.ts new file mode 100644 index 000000000..57014ced0 --- /dev/null +++ b/src/testing/module-mock-stack.ts @@ -0,0 +1,61 @@ +import { pathToFileURL } from "node:url" + +function normalizeStackPath(rawPath: string): string { + return rawPath.startsWith("file://") ? rawPath : pathToFileURL(rawPath).href +} + +function getStackCandidatePath(line: string): string | undefined { + const match = line.match(/(?:\()?(file:\/\/[^\s)]+|\/[^\s):]+):(\d+):(\d+)/) + return match?.[1] +} + +function isIgnoredCallerPath(candidatePath: string): boolean { + return ( + candidatePath.includes("/test-setup.ts") || + candidatePath.includes("/src/testing/module-mock-lifecycle.ts") || + candidatePath.includes("/src/testing/module-mock-stack.ts") + ) +} + +export function defaultGetCallerStack(): string { + return new Error().stack ?? "" +} + +export function resolveCallerUrlFromStack(stack: string): string { + const lines = stack.split("\n") + + for (const line of lines) { + const candidatePath = getStackCandidatePath(line) + if (!candidatePath) { + continue + } + + if (isIgnoredCallerPath(candidatePath)) { + continue + } + + return normalizeStackPath(candidatePath) + } + + return import.meta.url +} + +export function isModuleEvaluationStack(stack: string): boolean { + const lines = stack.split("\n") + + for (const [index, line] of lines.entries()) { + const candidatePath = getStackCandidatePath(line) + if (!candidatePath || isIgnoredCallerPath(candidatePath)) { + continue + } + + const nextFrame = lines[index + 1] ?? "" + return ( + nextFrame.includes("moduleEvaluation") || + nextFrame.includes("asyncModuleEvaluation") || + nextFrame.includes("loadAndEvaluateModule") + ) + } + + return false +} diff --git a/test-setup.ts b/test-setup.ts index 223f9f30e..5ac53c298 100644 --- a/test-setup.ts +++ b/test-setup.ts @@ -10,7 +10,9 @@ import { getOmoOpenCodeCacheDir } from "./src/shared/data-path" import { releaseAllPromptAsyncReservationsForTesting } from "./src/shared/prompt-async-gate" import { installModuleMockLifecycle } from "./src/testing/module-mock-lifecycle" -const { restoreModuleMocks } = installModuleMockLifecycle(mock) +const { beginTestMockTracking, endTestMockTracking, restoreModuleMocks } = installModuleMockLifecycle(mock, { + trackOnlyDuringActiveTest: true, +}) let environmentSnapshot: NodeJS.ProcessEnv = { ...process.env } let workingDirectorySnapshot = process.cwd() @@ -23,6 +25,7 @@ function cleanupRulesInjectorStorage(): void { } beforeEach(() => { + beginTestMockTracking() environmentSnapshot = { ...process.env } workingDirectorySnapshot = process.cwd() process.env.OMO_DISABLE_POSTHOG = "true" @@ -65,4 +68,5 @@ afterEach(() => { releaseAllPromptAsyncReservationsForTesting() mock.restore() restoreModuleMocks() + endTestMockTracking() })