fix(testing): preserve owner-scoped module mocks

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

@@ -156,6 +156,259 @@ describe("installModuleMockLifecycle active-test tracking", () => {
     ])
   })
 
+  test("treats single-frame parallel owner module calls as persistent while a different test is active", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/first.test.ts:5:1\n    at test (native:1:11)"
+    let callerUrl = "file:///repo/tests/first.test.ts"
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    beginTestMockTracking()
+    callerStack = "Error\n    at <anonymous> (/repo/tests/second.test.ts:5:1)"
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.module("./dependency", () => ({ named: "single-frame second" }))
+
+    // 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()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:single-frame second",
+      "delegate:restore",
+      "module:./dependency:single-frame second",
+      "module:resolved:file:///repo/tests/second.test.ts:./dependency:single-frame second",
+    ])
+  })
+
+  test("keeps beforeAll mocks persistent across owner test cleanup and restores them at afterAll", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/first.test.ts:5:1\n    at test (native:1:11)"
+    let callerUrl = "file:///repo/tests/first.test.ts"
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    beginTestMockTracking()
+    callerStack = "Error\n    at file:///repo/tests/second.test.ts:5:1\n    at beforeAll (native:1:11)"
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.module("./dependency", () => ({ named: "beforeAll second" }))
+
+    // 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()
+    endTestMockTracking()
+
+    callerStack = "Error\n    at file:///repo/tests/second.test.ts:20:1\n    at test (native:1:11)"
+    callerUrl = "file:///repo/tests/second.test.ts"
+    beginTestMockTracking()
+    callerStack = "Error\n    at file:///repo/tests/second.test.ts:25:1\n    at cleanup (native:1:11)"
+    mockApi.restore()
+    endTestMockTracking()
+
+    callerStack = "Error\n    at file:///repo/tests/second.test.ts:30:1\n    at afterAll (native:1:11)"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:beforeAll second",
+      "delegate:restore",
+      "module:./dependency:beforeAll second",
+      "module:resolved:file:///repo/tests/second.test.ts:./dependency:beforeAll second",
+      "delegate:restore",
+      "module:./dependency:beforeAll second",
+      "module:resolved:file:///repo/tests/second.test.ts:./dependency:beforeAll second",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:file:///repo/tests/second.test.ts:./dependency:original",
+    ])
+  })
+
+  test("clears only afterAll owner mocks while another test is active", () => {
+    // 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<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => 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" }))
+
+    callerStack = "Error\n    at file:///repo/tests/third.test.ts:10:1\n    at test (native:1:11)"
+    callerUrl = "file:///repo/tests/third.test.ts"
+    beginTestMockTracking()
+    callerStack = "Error\n    at file:///repo/tests/third.test.ts:15:1\n    at cleanup (native:1:11)"
+    mockApi.restore()
+
+    // when
+    callerStack = "Error\n    at file:///repo/tests/first.test.ts:20:1\n    at afterAll (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:./first:first top-level",
+      "module:resolved:file:///repo/tests/first.test.ts:./first:first top-level",
+      "module:./second:second top-level",
+      "module:resolved:file:///repo/tests/second.test.ts:./second:second top-level",
+      "delegate:restore",
+      "module:./first:original",
+      "module:resolved:file:///repo/tests/first.test.ts:./first:original",
+      "module:./second:second top-level",
+      "module:resolved:file:///repo/tests/second.test.ts:./second:second top-level",
+    ])
+  })
+
+  test("finalizes single-frame parallel owner restore while another test stays active", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/first.test.ts:5:1\n    at test (native:1:11)"
+    let callerUrl = "file:///repo/tests/first.test.ts"
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    beginTestMockTracking()
+    callerStack = "Error\n    at <anonymous> (/repo/tests/second.test.ts:5:1)"
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.module("./dependency", () => ({ named: "single-frame second" }))
+
+    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()
+
+    // when
+    callerStack = "Error\n    at <anonymous> (/repo/tests/second.test.ts:20:1)"
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:single-frame second",
+      "delegate:restore",
+      "module:./dependency:single-frame second",
+      "module:resolved:file:///repo/tests/second.test.ts:./dependency:single-frame second",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:file:///repo/tests/second.test.ts:./dependency:original",
+    ])
+  })
+
+  test("does not clear another owner's persistent mocks when unrelated afterAll runs during an active test", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/third.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    let callerUrl = "file:///repo/tests/third.test.ts"
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    const { beginTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => callerStack,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./third", () => ({ named: "third top-level" }))
+    callerStack = "Error\n    at file:///repo/tests/third.test.ts:10:1\n    at test (native:1:11)"
+    beginTestMockTracking()
+    callerStack = "Error\n    at file:///repo/tests/third.test.ts:15:1\n    at cleanup (native:1:11)"
+    mockApi.restore()
+
+    // when
+    callerStack = "Error\n    at <anonymous> (/repo/tests/unrelated.test.ts:20:1)"
+    callerUrl = "file:///repo/tests/unrelated.test.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./third:third top-level",
+      "delegate:restore",
+      "module:./third:third top-level",
+      "module:resolved:file:///repo/tests/third.test.ts:./third:third top-level",
+      "delegate:restore",
+      "module:./third:third top-level",
+      "module:resolved:file:///repo/tests/third.test.ts:./third:third top-level",
+    ])
+  })
+
   test("clears persistent module-evaluation snapshots when restore runs while inactive", () => {
     // given
     const events: string[] = []
@@ -352,6 +605,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
     // given
     const events: string[] = []
     let callerUrl = "file:///repo/tests/first.test.ts"
+    let activeOwnerInStackUrl = callerUrl
     const mockApi = {
       module: (specifier: string, factory: () => Record<string, unknown>) => {
         events.push(`module:${specifier}:${String(factory().named)}`)
@@ -362,7 +616,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
     }
 
     const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
-      getCallerStack: () => `Error\n    at ${callerUrl}:5:1\n    at cleanup (native:1:11)`,
+      getCallerStack: () => `Error\n    at ${callerUrl}:5:1\n    at ${activeOwnerInStackUrl}:6:1\n    at cleanup (native:1:11)`,
       getCallerUrl: () => callerUrl,
       trackOnlyDuringActiveTest: true,
       resolveSpecifier: (specifier) => `resolved:${specifier}`,
@@ -377,6 +631,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
 
     // when
     callerUrl = "file:///repo/tests/second.test.ts"
+    activeOwnerInStackUrl = callerUrl
     beginTestMockTracking()
     endTestMockTracking()
     callerUrl = "file:///repo/testing/test-setup.ts"

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

@@ -149,6 +149,50 @@ describe("installModuleMockLifecycle", () => {
     ])
   })
 
+  test("keeps other active owner mocks when one active owner restores", () => {
+    // given
+    const events: string[] = []
+    let callerUrl = "file:///repo/tests/first.test.ts"
+    const mockApi = {
+      module: (specifier: string, factory: () => Record<string, unknown>) => {
+        events.push(`module:${specifier}:${String(factory().named)}`)
+      },
+      restore: mock(() => {
+        events.push("delegate:restore")
+      }),
+    }
+
+    installModuleMockLifecycle(mockApi, {
+      getCallerUrl: () => callerUrl,
+      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./first", () => ({ named: "first mock" }))
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.module("./second", () => ({ named: "second mock" }))
+
+    // when
+    callerUrl = "file:///repo/tests/first.test.ts"
+    mockApi.restore()
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./first:first mock",
+      "module:./second:second mock",
+      "delegate:restore",
+      "module:./first:original",
+      "module:resolved:file:///repo/tests/first.test.ts:./first:original",
+      "module:./second:second mock",
+      "module:resolved:file:///repo/tests/second.test.ts:./second:second mock",
+      "delegate:restore",
+      "module:./second:original",
+      "module:resolved:file:///repo/tests/second.test.ts:./second:original",
+    ])
+  })
+
   test("captures the original module only once per resolved specifier", () => {
     // given
     let loadCount = 0

src/testing/module-mock-lifecycle.ts

@@ -1,6 +1,9 @@
 import { createRequire } from "node:module"
 import { fileURLToPath } from "node:url"
-import { defaultGetCallerStack, isModuleEvaluationStack, resolveCallerUrlFromStack } from "./module-mock-stack"
+import {
+  defaultGetCallerStack,
+  resolveCallerUrlFromStack,
+} from "./module-mock-stack"
 import { createRestoreExports } from "./module-mock-restore-exports"
 
 type MockModuleFactory = () => Record<string, unknown>
@@ -15,6 +18,7 @@ type ModuleLoadResult =
   | { ok: false; error: Error }
 
 type ModuleSnapshot = {
+  mockFactory: MockModuleFactory
   restoreSpecifiers: Set<string>
   restoreFactory: MockModuleFactory
 }
@@ -93,7 +97,7 @@ export function installModuleMockLifecycle(
   endTestMockTracking: () => void
   restoreModuleMocks: () => void
 } {
-  const snapshots = new Map<string, ModuleSnapshot>()
+  const snapshotsByOwner = new Map<string, Map<string, ModuleSnapshot>>()
   const persistentSnapshots = new Map<string, Map<string, PersistentModuleSnapshot>>()
   let lastRestoredSnapshots: ModuleSnapshot[] = []
   let lastRestoredSnapshotOwnerUrl: string | null = null
@@ -112,10 +116,55 @@ export function installModuleMockLifecycle(
     return options.getCallerUrl?.() ?? resolveCallerUrlFromStack(callerStack)
   }
 
-  function restoreModuleMocksForRestoreCall(callerUrl: string): void {
+  function hasActiveModuleMockOwner(ownerUrl: string): boolean {
+    return snapshotsByOwner.has(ownerUrl)
+  }
+
+  function stackReferencesOwnerUrl(stack: string, ownerUrl: string): boolean {
+    if (stack.includes(ownerUrl)) {
+      return true
+    }
+
+    if (!ownerUrl.startsWith("file://")) {
+      return false
+    }
+
+    if (!URL.canParse(ownerUrl)) {
+      return false
+    }
+
+    return stack.includes(fileURLToPath(ownerUrl))
+  }
+
+  function isParallelOwnerCall(callerStack: string, callerUrl: string): boolean {
+    return (
+      isActiveTest &&
+      options.trackOnlyDuringActiveTest === true &&
+      activeTestOwnerUrl !== null &&
+      callerUrl !== activeTestOwnerUrl &&
+      !stackReferencesOwnerUrl(callerStack, activeTestOwnerUrl)
+    )
+  }
+
+  function resolveActiveRestoreOwner(callerUrl: string): string | null {
+    if (hasActiveModuleMockOwner(callerUrl)) {
+      return callerUrl
+    }
+
+    if (activeTestOwnerUrl && hasActiveModuleMockOwner(activeTestOwnerUrl)) {
+      return activeTestOwnerUrl
+    }
+
+    return options.trackOnlyDuringActiveTest === true ? callerUrl : null
+  }
+
+  function restoreModuleMocksForRestoreCall(callerUrl: string, restoreOwnerUrl: string | null = callerUrl): void {
+    const callerSnapshots = restoreOwnerUrl ? snapshotsByOwner.get(restoreOwnerUrl) : undefined
     const snapshotsToRestore =
-      snapshots.size > 0
+      restoreOwnerUrl === null && snapshotsByOwner.size > 0
-        ? Array.from(snapshots.values())
+        ? Array.from(snapshotsByOwner.values()).flatMap((snapshots) => Array.from(snapshots.values()))
+        : callerSnapshots && callerSnapshots.size > 0
+        ? Array.from(callerSnapshots.values())
         : callerUrl === lastRestoredSnapshotOwnerUrl
           ? lastRestoredSnapshots
           : []
@@ -126,10 +175,26 @@ export function installModuleMockLifecycle(
       }
     }
 
-    if (snapshots.size > 0) {
+    if (restoreOwnerUrl === null && snapshotsByOwner.size > 0) {
       lastRestoredSnapshots = snapshotsToRestore
       lastRestoredSnapshotOwnerUrl = callerUrl
-      snapshots.clear()
+      snapshotsByOwner.clear()
+    } else if (callerSnapshots && callerSnapshots.size > 0 && restoreOwnerUrl) {
+      lastRestoredSnapshots = snapshotsToRestore
+      lastRestoredSnapshotOwnerUrl = callerUrl
+      snapshotsByOwner.delete(restoreOwnerUrl)
+    }
+
+    for (const [ownerUrl, snapshots] of snapshotsByOwner) {
+      if (restoreOwnerUrl === null || ownerUrl === restoreOwnerUrl) {
+        continue
+      }
+
+      for (const snapshot of snapshots.values()) {
+        for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+          delegateModule(restoreSpecifier, snapshot.mockFactory)
+        }
+      }
     }
   }
 
@@ -186,8 +251,12 @@ export function installModuleMockLifecycle(
     return false
   }
 
+  function hasModuleMockOwner(ownerUrl: string): boolean {
+    return hasActiveModuleMockOwner(ownerUrl) || hasPersistentModuleMockOwner(ownerUrl)
+  }
+
   function resolveInactiveRestoreOwner(callerUrl: string): string {
-    if (hasPersistentModuleMockOwner(callerUrl)) {
+    if (hasModuleMockOwner(callerUrl)) {
       return callerUrl
     }
 
@@ -195,7 +264,7 @@ export function installModuleMockLifecycle(
   }
 
   function restoreModuleMocks(): void {
-    if (snapshots.size === 0) {
+    if (snapshotsByOwner.size === 0) {
       return
     }
 
@@ -221,14 +290,9 @@ export function installModuleMockLifecycle(
     lastRestoredSnapshots = []
     const callerStack = getCallerStack()
     const callerUrl = getCallerUrl(callerStack)
-    const isParallelFileEvaluationMock =
+    const isParallelOwnerModuleCall = isParallelOwnerCall(callerStack, callerUrl)
-      isActiveTest &&
-      options.trackOnlyDuringActiveTest === true &&
-      activeTestOwnerUrl !== null &&
-      callerUrl !== activeTestOwnerUrl &&
-      isModuleEvaluationStack(callerStack)
 
-    if ((!isActiveTest || isParallelFileEvaluationMock) && isPersistentModuleMockOwner(callerUrl)) {
+    if ((!isActiveTest || isParallelOwnerModuleCall) && isPersistentModuleMockOwner(callerUrl)) {
       const resolvedSpecifier = resolveSpecifier(specifier, callerUrl)
       const snapshotsByOwner = persistentSnapshots.get(resolvedSpecifier) ?? new Map<string, PersistentModuleSnapshot>()
       const existingSnapshot = snapshotsByOwner.get(callerUrl)
@@ -250,20 +314,24 @@ export function installModuleMockLifecycle(
 
     if (isActiveTest) {
       const resolvedSpecifier = resolveSpecifier(specifier, callerUrl)
-      const existingSnapshot = snapshots.get(resolvedSpecifier)
+      const ownerSnapshots = snapshotsByOwner.get(callerUrl) ?? new Map<string, ModuleSnapshot>()
+      const existingSnapshot = ownerSnapshots.get(resolvedSpecifier)
 
       if (existingSnapshot) {
         existingSnapshot.restoreSpecifiers.add(specifier)
         existingSnapshot.restoreSpecifiers.add(resolvedSpecifier)
+        existingSnapshot.mockFactory = factory
       } else {
         const originalModule = loadOriginalModule(specifier, callerUrl)
 
         if (originalModule.ok) {
           const restoreExports = createRestoreExports(originalModule.value)
-          snapshots.set(resolvedSpecifier, {
+          ownerSnapshots.set(resolvedSpecifier, {
+            mockFactory: factory,
             restoreSpecifiers: new Set([specifier, resolvedSpecifier]),
             restoreFactory: () => restoreExports,
           })
+          snapshotsByOwner.set(callerUrl, ownerSnapshots)
         }
       }
     }
@@ -275,17 +343,27 @@ export function installModuleMockLifecycle(
     const callerStack = getCallerStack()
     const callerUrl = getCallerUrl(callerStack)
     const result = delegateRestore()
-    if (!isActiveTest) {
+    const isParallelOwnerRestoreCall = isParallelOwnerCall(callerStack, callerUrl)
-      restoreModuleMocksForRestoreCall(callerUrl)
+    const shouldTreatRestoreAsInactive =
-      snapshots.clear()
+      !isActiveTest || (isParallelOwnerRestoreCall && hasModuleMockOwner(callerUrl))
+
+    if (shouldTreatRestoreAsInactive) {
+      const ownerUrl = resolveInactiveRestoreOwner(callerUrl)
+      restoreModuleMocksForRestoreCall(callerUrl, ownerUrl)
       lastRestoredSnapshots = []
       lastRestoredSnapshotOwnerUrl = null
-      clearPersistentModuleMocksForOwner(resolveInactiveRestoreOwner(callerUrl), hasStartedTest)
+      clearPersistentModuleMocksForOwner(ownerUrl, hasStartedTest)
       restorePersistentModuleMocksForRestoreCall(false)
       return result
     }
 
-    restoreModuleMocksForRestoreCall(callerUrl)
+    if (isParallelOwnerRestoreCall) {
+      restoreModuleMocksForRestoreCall(callerUrl, callerUrl)
+      restorePersistentModuleMocksForRestoreCall(true)
+      return result
+    }
+
+    restoreModuleMocksForRestoreCall(callerUrl, resolveActiveRestoreOwner(callerUrl))
     restorePersistentModuleMocksForRestoreCall(true)
     return result
   }