fix(testing): scope module mock restore snapshots

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

@@ -0,0 +1,57 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, mock, test } from "bun:test"
+import { installModuleMockLifecycle } from "./module-mock-lifecycle"
+
+describe("installModuleMockLifecycle inactive restore ownership", () => {
+  test("keeps unrelated re-applied persistent mocks when inactive restore caller is unresolved", () => {
+    // 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")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => `Error\n    at ${callerUrl}:5:1`,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      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" }))
+    callerUrl = "file:///repo/tests/first.test.ts"
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // when
+    callerUrl = "file:///repo/testing/unknown-owner.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./first:first mock",
+      "module:./second:second mock",
+      "delegate:restore",
+      "module:./first:first mock",
+      "module:resolved:file:///repo/tests/first.test.ts:./first:first mock",
+      "module:./second:second mock",
+      "module:resolved:file:///repo/tests/second.test.ts:./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",
+    ])
+  })
+})

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

@@ -311,7 +311,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
     expect(loadOriginalModule).toHaveBeenCalledTimes(1)
   })
 
-  test("reapplies the last active restore snapshot when a later inactive restore runs", () => {
+  test("does not reapply the last active restore snapshot after tracking ends", () => {
     // given
     const events: string[] = []
     const mockApi = {
@@ -345,8 +345,50 @@ describe("installModuleMockLifecycle active-test tracking", () => {
       "module:./dependency:original",
       "module:resolved:./dependency:original",
       "delegate:restore",
+    ])
+  })
+
+  test("does not replay a prior test snapshot during the next test cleanup", () => {
+    // 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")
+      }),
+    }
+
+    const { beginTestMockTracking, endTestMockTracking } = installModuleMockLifecycle(mockApi, {
+      getCallerStack: () => `Error\n    at ${callerUrl}:5:1\n    at cleanup (native:1:11)`,
+      getCallerUrl: () => callerUrl,
+      trackOnlyDuringActiveTest: true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    beginTestMockTracking()
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    callerUrl = "file:///repo/testing/test-setup.ts"
+    mockApi.restore()
+    endTestMockTracking()
+
+    // when
+    callerUrl = "file:///repo/tests/second.test.ts"
+    beginTestMockTracking()
+    endTestMockTracking()
+    callerUrl = "file:///repo/testing/test-setup.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
       "module:./dependency:original",
       "module:resolved:./dependency:original",
+      "delegate:restore",
     ])
   })
 

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

@@ -114,6 +114,41 @@ describe("installModuleMockLifecycle", () => {
     ])
   })
 
+  test("does not replay the last restore snapshot for a different restore caller", () => {
+    // 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) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    mockApi.restore()
+    callerUrl = "file:///repo/tests/second.test.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+      "delegate:restore",
+    ])
+  })
+
   test("captures the original module only once per resolved specifier", () => {
     // given
     let loadCount = 0

src/testing/module-mock-lifecycle.ts

@@ -96,9 +96,11 @@ export function installModuleMockLifecycle(
   const snapshots = new Map<string, ModuleSnapshot>()
   const persistentSnapshots = new Map<string, Map<string, PersistentModuleSnapshot>>()
   let lastRestoredSnapshots: ModuleSnapshot[] = []
+  let lastRestoredSnapshotOwnerUrl: string | null = null
   let isActiveTest = !options.trackOnlyDuringActiveTest
   let hasStartedTest = false
   let activeTestOwnerUrl: string | null = null
+  let lastActiveTestOwnerUrl: string | null = null
   const delegateModule = mockApi.module.bind(mockApi)
   const delegateRestore = mockApi.restore.bind(mockApi)
   const getCallerStack = options.getCallerStack ?? defaultGetCallerStack
@@ -110,8 +112,13 @@ export function installModuleMockLifecycle(
     return options.getCallerUrl?.() ?? resolveCallerUrlFromStack(callerStack)
   }
 
-  function restoreModuleMocksForRestoreCall(): void {
-    const snapshotsToRestore = snapshots.size > 0 ? Array.from(snapshots.values()) : lastRestoredSnapshots
+  function restoreModuleMocksForRestoreCall(callerUrl: string): void {
+    const snapshotsToRestore =
+      snapshots.size > 0
+        ? Array.from(snapshots.values())
+        : callerUrl === lastRestoredSnapshotOwnerUrl
+          ? lastRestoredSnapshots
+          : []
 
     for (const snapshot of snapshotsToRestore) {
       for (const restoreSpecifier of snapshot.restoreSpecifiers) {
@@ -121,6 +128,7 @@ export function installModuleMockLifecycle(
 
     if (snapshots.size > 0) {
       lastRestoredSnapshots = snapshotsToRestore
+      lastRestoredSnapshotOwnerUrl = callerUrl
       snapshots.clear()
     }
   }
@@ -153,41 +161,37 @@ export function installModuleMockLifecycle(
   function clearPersistentModuleMocksForOwner(
     ownerUrl: string,
     restoreOriginals: boolean,
-    forceRestoreOriginals = false,
   ): void {
-    let clearedOwnerSnapshot = false
-
     for (const [resolvedSpecifier, snapshotsByOwner] of persistentSnapshots) {
       const snapshot = snapshotsByOwner.get(ownerUrl)
       if (snapshot) {
-        if (restoreOriginals && (forceRestoreOriginals || snapshot.reappliedDuringActiveRestore)) {
+        if (restoreOriginals && snapshot.reappliedDuringActiveRestore) {
           restorePersistentOriginals(snapshot, ownerUrl)
         }
         snapshotsByOwner.delete(ownerUrl)
-        clearedOwnerSnapshot = true
       }
       if (snapshotsByOwner.size === 0) {
         persistentSnapshots.delete(resolvedSpecifier)
       }
     }
+  }
 
-    if (clearedOwnerSnapshot || !restoreOriginals) {
-      return
-    }
-
-    for (const [resolvedSpecifier, snapshotsByOwner] of persistentSnapshots) {
-      for (const [snapshotOwnerUrl, snapshot] of snapshotsByOwner) {
-        if (!snapshot.reappliedDuringActiveRestore) {
-          continue
-        }
-
-        restorePersistentOriginals(snapshot, snapshotOwnerUrl)
-        snapshotsByOwner.delete(snapshotOwnerUrl)
-      }
-      if (snapshotsByOwner.size === 0) {
-        persistentSnapshots.delete(resolvedSpecifier)
+  function hasPersistentModuleMockOwner(ownerUrl: string): boolean {
+    for (const snapshotsByOwner of persistentSnapshots.values()) {
+      if (snapshotsByOwner.has(ownerUrl)) {
+        return true
       }
     }
+
+    return false
+  }
+
+  function resolveInactiveRestoreOwner(callerUrl: string): string {
+    if (hasPersistentModuleMockOwner(callerUrl)) {
+      return callerUrl
+    }
+
+    return lastActiveTestOwnerUrl ?? callerUrl
   }
 
   function restoreModuleMocks(): void {
@@ -195,7 +199,7 @@ export function installModuleMockLifecycle(
       return
     }
 
-    restoreModuleMocksForRestoreCall()
+    restoreModuleMocksForRestoreCall(getCallerUrl(getCallerStack()))
   }
 
   function beginTestMockTracking(): void {
@@ -203,11 +207,14 @@ export function installModuleMockLifecycle(
     isActiveTest = true
     const callerStack = getCallerStack()
     activeTestOwnerUrl = getCallerUrl(callerStack)
+    lastActiveTestOwnerUrl = activeTestOwnerUrl
   }
 
   function endTestMockTracking(): void {
     isActiveTest = !options.trackOnlyDuringActiveTest
     activeTestOwnerUrl = null
+    lastRestoredSnapshots = []
+    lastRestoredSnapshotOwnerUrl = null
   }
 
   mockApi.module = (specifier: string, factory: MockModuleFactory): unknown => {
@@ -269,15 +276,16 @@ export function installModuleMockLifecycle(
     const callerUrl = getCallerUrl(callerStack)
     const result = delegateRestore()
     if (!isActiveTest) {
-      restoreModuleMocksForRestoreCall()
+      restoreModuleMocksForRestoreCall(callerUrl)
       snapshots.clear()
       lastRestoredSnapshots = []
-      clearPersistentModuleMocksForOwner(callerUrl, hasStartedTest)
+      lastRestoredSnapshotOwnerUrl = null
+      clearPersistentModuleMocksForOwner(resolveInactiveRestoreOwner(callerUrl), hasStartedTest)
       restorePersistentModuleMocksForRestoreCall(false)
       return result
     }
 
-    restoreModuleMocksForRestoreCall()
+    restoreModuleMocksForRestoreCall(callerUrl)
     restorePersistentModuleMocksForRestoreCall(true)
     return result
   }