fix(testing): restore mock specifier aliases

2026-05-30 15:34:11 +09:00
parent 78cd34b17b
commit dadcdbabd5
4 changed files with 414 additions and 39 deletions
@@ -42,7 +42,7 @@ describe("installModuleMockLifecycle Bun integration", () => {
    mockApi.restore()

    // then
-    const restoreCall = moduleCalls.at(1)
+    const restoreCall = moduleCalls.find((call) => /\/src\/testing\/module-mock-lifecycle-fixture\.ts$/.test(call.specifier))
    expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-fixture\.ts$/)
    expect(restoreCall?.value.named).toBe("original")
  })
@@ -64,7 +64,7 @@ describe("installModuleMockLifecycle Bun integration", () => {
    mockApi.restore()

    // then
-    const restoreCall = moduleCalls.at(1)
+    const restoreCall = moduleCalls.find((call) => /\/src\/testing\/module-mock-lifecycle-nested\/fixture\.ts$/.test(call.specifier))
    expect(restoreCall?.specifier).toMatch(/\/src\/testing\/module-mock-lifecycle-nested\/fixture\.ts$/)
    expect(restoreCall?.value.named).toBe("original")
  })
@@ -21,6 +21,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
      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,
+      isPersistentModuleMockOwner: () => true,
      resolveSpecifier: (specifier) => `resolved:${specifier}`,
      loadOriginalModule,
    })
@@ -76,6 +77,42 @@ describe("installModuleMockLifecycle active-test tracking", () => {
    ])
  })

+  test("restores active in-test mocks even when Bun stack includes module evaluation", () => {
+    // given
+    const events: string[] = []
+    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 file:///repo/tests/example.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/example.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => 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[] = []
@@ -92,6 +129,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
      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,
+      isPersistentModuleMockOwner: () => true,
      resolveSpecifier: (specifier) => `resolved:${specifier}`,
      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
    })
@@ -111,6 +149,164 @@ describe("installModuleMockLifecycle active-test tracking", () => {
    ])
  })

+  test("does not restore ordinary inactive mocks before their owner test cleanup runs", () => {
+    // given
+    const events: string[] = []
+    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 file:///repo/tests/late-loaded.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/late-loaded.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    beginTestMockTracking()
+    endTestMockTracking()
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+    ])
+  })
+
+  test("restores original exports for beforeAll-style inactive mocks after owner test cleanup", () => {
+    // given
+    const events: string[] = []
+    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 file:///repo/tests/before-all.test.ts:5:1\n    at beforeAll (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/before-all.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+  })
+
+  test("restores original exports for persistent module-evaluation mocks after owner test cleanup", () => {
+    // given
+    const events: string[] = []
+    const loadOriginalModule = mock(() => ({ ok: true as const, value: { named: "original" } }))
+    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 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,
+      isPersistentModuleMockOwner: () => true,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule,
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+    expect(loadOriginalModule).toHaveBeenCalledTimes(1)
+  })
+
+  test("reapplies the last active restore snapshot when a later inactive restore runs", () => {
+    // given
+    const events: string[] = []
+    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 file:///repo/tests/example.test.ts:5:1\n    at cleanup (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()
+    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",
+    ])
+  })
+
  test("keeps unrelated module-evaluation mocks after inactive restore", () => {
    // given
    const events: string[] = []
@@ -129,6 +325,7 @@ describe("installModuleMockLifecycle active-test tracking", () => {
      getCallerStack: () => callerStack,
      getCallerUrl: () => callerUrl,
      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => true,
      resolveSpecifier: (specifier, ownerUrl) => `resolved:${ownerUrl}:${specifier}`,
      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
    })
@@ -152,4 +349,128 @@ describe("installModuleMockLifecycle active-test tracking", () => {
      "module:resolved:file:///repo/tests/second.test.ts:./second:second top-level",
    ])
  })
+
+  test("keeps module-evaluation mocks when inactive restore cannot resolve the owner", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/top-level.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    let callerUrl = "file:///repo/tests/top-level.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) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+
+    // when
+    callerStack = "Error\n    at nativeAfterAll (native:1:11)"
+    callerUrl = "file:///repo/testing/unknown-owner.ts"
+    mockApi.restore()
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+    ])
+  })
+
+  test("restores re-applied persistent mocks when inactive restore cannot resolve the owner", () => {
+    // given
+    const events: string[] = []
+    let callerStack = "Error\n    at file:///repo/tests/top-level.test.ts:5:1\n    at moduleEvaluation (native:1:11)"
+    let callerUrl = "file:///repo/tests/top-level.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) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    mockApi.module("./dependency", () => ({ named: "mocked" }))
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // when
+    callerStack = "Error\n    at nativeAfterAll (native:1:11)"
+    callerUrl = "file:///repo/testing/unknown-owner.ts"
+    mockApi.restore()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:mocked",
+      "module:resolved:./dependency:mocked",
+      "delegate:restore",
+      "module:./dependency:original",
+      "module:resolved:./dependency:original",
+    ])
+  })
+
+  test("does not reapply ordinary module-evaluation mocks after active restore", () => {
+    // given
+    const events: string[] = []
+    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 file:///repo/tests/file-local.test.ts:5:1\n    at moduleEvaluation (native:1:11)",
+      getCallerUrl: () => "file:///repo/tests/file-local.test.ts",
+      trackOnlyDuringActiveTest: true,
+      isPersistentModuleMockOwner: () => false,
+      resolveSpecifier: (specifier) => `resolved:${specifier}`,
+      loadOriginalModule: () => ({ ok: true, value: { named: "original" } }),
+    })
+
+    // when
+    mockApi.module("./dependency", () => ({ named: "file local" }))
+    beginTestMockTracking()
+    mockApi.restore()
+    endTestMockTracking()
+
+    // then
+    expect(events).toEqual([
+      "module:./dependency:file local",
+      "delegate:restore",
+    ])
+  })
 })
@@ -4,7 +4,7 @@ import { describe, expect, mock, test } from "bun:test"
 import { installModuleMockLifecycle } from "./module-mock-lifecycle"

 describe("installModuleMockLifecycle", () => {
-  test("preserves plain export identity while cloning ESM namespace snapshots", () => {
+  test("restores every mock specifier while cloning ESM namespace snapshots", () => {
    // given
    const originalExports = { named: "original" }
    const moduleNamespaceExports = Object.defineProperty({ named: "namespace" }, Symbol.toStringTag, {
@@ -34,12 +34,17 @@ describe("installModuleMockLifecycle", () => {

    // then
    const plainRestoreCall = moduleCalls.find((call) => call.specifier === "./dependency" && call.value.named === "original")
-    const namespaceRestoreCall = moduleCalls.find(
+    const namespaceOriginalRestoreCall = moduleCalls.find(
+      (call) => call.specifier === "./namespace" && call.value.named === "namespace",
+    )
+    const namespaceResolvedRestoreCall = 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)
+    expect(namespaceOriginalRestoreCall?.value).toEqual({ named: "namespace" })
+    expect(namespaceOriginalRestoreCall?.value).not.toBe(moduleNamespaceExports)
+    expect(namespaceResolvedRestoreCall?.value).toEqual({ named: "namespace" })
+    expect(namespaceResolvedRestoreCall?.value).not.toBe(moduleNamespaceExports)
  })

  test("clears tracked snapshots after the delegate restore runs", () => {
@@ -1,6 +1,6 @@
 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"

 type MockModuleFactory = () => Record<string, unknown>

@@ -14,13 +14,13 @@ type ModuleLoadResult =
  | { ok: false; error: Error }

 type ModuleSnapshot = {
-  restoreOriginalSpecifiers: boolean
  restoreSpecifiers: Set<string>
  restoreFactory: MockModuleFactory
 }

 type PersistentModuleSnapshot = {
-  ownerUrls: Set<string>
+  originalSpecifier: string
+  reappliedDuringActiveRestore: boolean
  restoreSpecifiers: Set<string>
  restoreFactory: MockModuleFactory
 }
@@ -29,6 +29,7 @@ type ModuleMockLifecycleOptions = {
  getCallerStack?: () => string
  getCallerUrl?: () => string
  trackOnlyDuringActiveTest?: boolean
+  isPersistentModuleMockOwner?: (callerUrl: string) => boolean
  resolveSpecifier?: (specifier: string, callerUrl: string) => string
  loadOriginalModule?: (specifier: string, callerUrl: string) => ModuleLoadResult
 }
@@ -47,10 +48,6 @@ function isModuleNamespaceObject(moduleValue: Record<string, unknown>): boolean
  return Object.prototype.toString.call(moduleValue) === "[object Module]"
 }

-function shouldRestoreOriginalSpecifier(moduleValue: unknown): boolean {
-  return !isModuleExports(moduleValue) || !isModuleNamespaceObject(moduleValue)
-}
-
 function createRestoreExports(moduleValue: unknown): Record<string, unknown> {
  if (typeof moduleValue === "function") {
    const functionExports = Object.assign({}, moduleValue)
@@ -111,6 +108,10 @@ function defaultLoadOriginalModule(specifier: string, callerUrl: string): Module
  }
 }

+function defaultIsPersistentModuleMockOwner(_callerUrl: string): boolean {
+  return true
+}
+
 export function installModuleMockLifecycle(
  mockApi: MockApi,
  options: ModuleMockLifecycleOptions = {},
@@ -120,14 +121,16 @@ export function installModuleMockLifecycle(
  restoreModuleMocks: () => void
 } {
  const snapshots = new Map<string, ModuleSnapshot>()
-  const persistentSnapshots = new Map<string, PersistentModuleSnapshot>()
+  const persistentSnapshots = new Map<string, Map<string, PersistentModuleSnapshot>>()
  let lastRestoredSnapshots: ModuleSnapshot[] = []
  let isActiveTest = !options.trackOnlyDuringActiveTest
+  let hasStartedTest = false
  const delegateModule = mockApi.module.bind(mockApi)
  const delegateRestore = mockApi.restore.bind(mockApi)
  const getCallerStack = options.getCallerStack ?? defaultGetCallerStack
  const resolveSpecifier = options.resolveSpecifier ?? defaultResolveSpecifier
  const loadOriginalModule = options.loadOriginalModule ?? defaultLoadOriginalModule
+  const isPersistentModuleMockOwner = options.isPersistentModuleMockOwner ?? defaultIsPersistentModuleMockOwner

  function getCallerUrl(callerStack: string): string {
    return options.getCallerUrl?.() ?? resolveCallerUrlFromStack(callerStack)
@@ -148,18 +151,66 @@ export function installModuleMockLifecycle(
    }
  }

-  function restorePersistentModuleMocksForRestoreCall(): void {
-    for (const snapshot of persistentSnapshots.values()) {
-      for (const restoreSpecifier of snapshot.restoreSpecifiers) {
-        delegateModule(restoreSpecifier, snapshot.restoreFactory)
+  function restorePersistentModuleMocksForRestoreCall(markReapplied: boolean): void {
+    for (const snapshotsByOwner of persistentSnapshots.values()) {
+      for (const snapshot of snapshotsByOwner.values()) {
+        if (markReapplied) {
+          snapshot.reappliedDuringActiveRestore = true
+        }
+        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) {
+  function restorePersistentOriginals(snapshot: PersistentModuleSnapshot, ownerUrl: string): void {
+    const originalModule = loadOriginalModule(snapshot.originalSpecifier, ownerUrl)
+    if (!originalModule.ok) {
+      return
+    }
+
+    const originalFactory = () => createRestoreExports(originalModule.value)
+    for (const restoreSpecifier of snapshot.restoreSpecifiers) {
+      delegateModule(restoreSpecifier, originalFactory)
+    }
+  }
+
+  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)) {
+          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)
      }
    }
@@ -174,6 +225,7 @@ export function installModuleMockLifecycle(
  }

  function beginTestMockTracking(): void {
+    hasStartedTest = true
    isActiveTest = true
  }

@@ -185,22 +237,24 @@ export function installModuleMockLifecycle(
    lastRestoredSnapshots = []
    const callerStack = getCallerStack()
    const callerUrl = getCallerUrl(callerStack)
-    const isModuleEvaluation = isModuleEvaluationStack(callerStack)

-    if (isModuleEvaluation) {
+    if (!isActiveTest && isPersistentModuleMockOwner(callerUrl)) {
      const resolvedSpecifier = resolveSpecifier(specifier, callerUrl)
-      const existingSnapshot = persistentSnapshots.get(resolvedSpecifier)
+      const snapshotsByOwner = persistentSnapshots.get(resolvedSpecifier) ?? new Map<string, PersistentModuleSnapshot>()
+      const existingSnapshot = snapshotsByOwner.get(callerUrl)
      if (existingSnapshot) {
-        existingSnapshot.ownerUrls.add(callerUrl)
        existingSnapshot.restoreSpecifiers.add(specifier)
        existingSnapshot.restoreSpecifiers.add(resolvedSpecifier)
+        existingSnapshot.restoreFactory = factory
      } else {
-        persistentSnapshots.set(resolvedSpecifier, {
-          ownerUrls: new Set([callerUrl]),
+        snapshotsByOwner.set(callerUrl, {
+          originalSpecifier: specifier,
+          reappliedDuringActiveRestore: false,
          restoreSpecifiers: new Set([specifier, resolvedSpecifier]),
          restoreFactory: factory,
        })
      }
+      persistentSnapshots.set(resolvedSpecifier, snapshotsByOwner)
      return delegateModule(specifier, factory)
    }

@@ -209,21 +263,15 @@ export function installModuleMockLifecycle(
      const existingSnapshot = snapshots.get(resolvedSpecifier)

      if (existingSnapshot) {
-        if (existingSnapshot.restoreOriginalSpecifiers) {
-          existingSnapshot.restoreSpecifiers.add(specifier)
-        }
+        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],
-            ),
+            restoreSpecifiers: new Set([specifier, resolvedSpecifier]),
            restoreFactory: () => restoreExports,
          })
        }
@@ -238,15 +286,16 @@ export function installModuleMockLifecycle(
    const callerUrl = getCallerUrl(callerStack)
    const result = delegateRestore()
    if (!isActiveTest) {
+      restoreModuleMocksForRestoreCall()
      snapshots.clear()
      lastRestoredSnapshots = []
-      clearPersistentModuleMocksForOwner(callerUrl)
-      restorePersistentModuleMocksForRestoreCall()
+      clearPersistentModuleMocksForOwner(callerUrl, hasStartedTest)
+      restorePersistentModuleMocksForRestoreCall(false)
      return result
    }

    restoreModuleMocksForRestoreCall()
-    restorePersistentModuleMocksForRestoreCall()
+    restorePersistentModuleMocksForRestoreCall(true)
    return result
  }