feat(testing): add module mock lifecycle utilities

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

src/testing/module-mock-lifecycle.test.ts

@@ -0,0 +1,82 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, mock, test } from "bun:test"
+import { installModuleMockLifecycle } from "./module-mock-lifecycle"
+
+describe("installModuleMockLifecycle", () => {
+  test("restores the original module exports on mock.restore", () => {
+    // given
+    const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        moduleCalls.push({ specifier, value: factory() })
+      },
+      restore: mock(() => {}),
+    }
+
+    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()
+
+    // then
+    expect(moduleCalls).toEqual([
+      { specifier: "./dependency", value: { named: "mocked" } },
+      { specifier: "resolved:./dependency", value: { named: "original" } },
+    ])
+  })
+
+  test("captures the original module only once per resolved specifier", () => {
+    // given
+    let loadCount = 0
+    const mockApi = {
+      module: mock(() => {}),
+      restore: mock(() => {}),
+    }
+
+    installModuleMockLifecycle(mockApi, {
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      resolveSpecifier: () => "file:///repo/src/dependency.ts",
+      loadOriginalModule: () => {
+        loadCount += 1
+        return { ok: true, value: { named: "original" } }
+      },
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "first" }))
+    mockApi.module("./dependency", () => ({ named: "second" }))
+
+    // then
+    expect(loadCount).toBe(1)
+  })
+
+  test("does not restore unresolved modules to avoid cleanup errors", () => {
+    // given
+    const moduleCalls: Array<{ specifier: string; value: Record<string, unknown> }> = []
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        moduleCalls.push({ specifier, value: factory() })
+      },
+      restore: mock(() => {}),
+    }
+
+    installModuleMockLifecycle(mockApi, {
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      resolveSpecifier: (specifier) => specifier,
+      loadOriginalModule: () => ({ ok: false, error: new Error("Cannot find module") }),
+    })
+
+    // when
+    mockApi.module("virtual:missing", () => ({ named: "mocked" }))
+    mockApi.restore()
+
+    // then - only the original mock call, no restore call for unresolved module
+    expect(moduleCalls).toEqual([{ specifier: "virtual:missing", value: { named: "mocked" } }])
+  })
+})

src/testing/module-mock-lifecycle.ts

@@ -0,0 +1,143 @@
+import { createRequire } from "node:module"
+import { pathToFileURL } from "node:url"
+
+type MockModuleFactory = () => Record<string, unknown>
+
+type MockApi = {
+  module: (specifier: string, factory: MockModuleFactory) => unknown
+  restore: () => unknown
+}
+
+type ModuleLoadResult =
+  | { ok: true; value: unknown }
+  | { ok: false; error: Error }
+
+type ModuleSnapshot = {
+  restoreSpecifier: string
+  restoreFactory: MockModuleFactory
+}
+
+type ModuleMockLifecycleOptions = {
+  getCallerUrl?: () => string
+  resolveSpecifier?: (specifier: string, callerUrl: string) => string
+  loadOriginalModule?: (specifier: string, callerUrl: string) => ModuleLoadResult
+}
+
+function toError(error: unknown): Error {
+  if (error instanceof Error) {
+    return error
+  }
+
+  return new Error(String(error))
+}
+
+function cloneModuleExports(moduleValue: unknown): Record<string, unknown> {
+  if (typeof moduleValue === "function") {
+    const functionExports = Object.assign({}, moduleValue)
+    return {
+      ...functionExports,
+      default: moduleValue,
+    }
+  }
+
+  if (moduleValue && typeof moduleValue === "object") {
+    return { ...(moduleValue as Record<string, unknown>) }
+  }
+
+  return { default: moduleValue }
+}
+
+function normalizeStackPath(rawPath: string): string {
+  if (rawPath.startsWith("file://")) {
+    return rawPath
+  }
+
+  return pathToFileURL(rawPath).href
+}
+
+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 defaultResolveSpecifier(specifier: string, callerUrl: string): string {
+  try {
+    return import.meta.resolve(specifier, callerUrl)
+  } catch {
+    return specifier
+  }
+}
+
+function defaultLoadOriginalModule(specifier: string, callerUrl: string): ModuleLoadResult {
+  try {
+    const require = createRequire(callerUrl)
+    return { ok: true, value: require(specifier) }
+  } catch (error) {
+    return { ok: false, error: toError(error) }
+  }
+}
+
+export function installModuleMockLifecycle(
+  mockApi: MockApi,
+  options: ModuleMockLifecycleOptions = {},
+): { restoreModuleMocks: () => void } {
+  const snapshots = new Map<string, ModuleSnapshot>()
+  const delegateModule = mockApi.module.bind(mockApi)
+  const delegateRestore = mockApi.restore.bind(mockApi)
+  const getCallerUrl = options.getCallerUrl ?? defaultGetCallerUrl
+  const resolveSpecifier = options.resolveSpecifier ?? defaultResolveSpecifier
+  const loadOriginalModule = options.loadOriginalModule ?? defaultLoadOriginalModule
+
+  function restoreModuleMocks(): void {
+    for (const snapshot of snapshots.values()) {
+      delegateModule(snapshot.restoreSpecifier, snapshot.restoreFactory)
+    }
+
+    snapshots.clear()
+  }
+
+  mockApi.module = (specifier: string, factory: MockModuleFactory): unknown => {
+    const callerUrl = getCallerUrl()
+    const restoreSpecifier = resolveSpecifier(specifier, callerUrl)
+
+    if (!snapshots.has(restoreSpecifier)) {
+      const originalModule = loadOriginalModule(specifier, callerUrl)
+
+      if (originalModule.ok) {
+        const clonedExports = cloneModuleExports(originalModule.value)
+        snapshots.set(restoreSpecifier, {
+          restoreSpecifier,
+          restoreFactory: () => ({ ...clonedExports }),
+        })
+      }
+    }
+
+    return delegateModule(specifier, factory)
+  }
+
+  mockApi.restore = (): unknown => {
+    restoreModuleMocks()
+    return delegateRestore()
+  }
+
+  return { restoreModuleMocks }
+}