bun:test: make mock.restore() revert mock.module() overrides

docs/test/mocks.mdx

@@ -451,7 +451,9 @@ test("resetting all mocks", () => {
 
 ### Restore All Mocks
 
-`mock.restore()` restores every mock at once, instead of calling `mockFn.mockRestore()` on each one. It does not reset modules overridden with `mock.module()`.
+`mock.restore()` restores every mock at once, instead of calling `mockFn.mockRestore()` on each one. It also undoes `mock.module()` calls made inside tests and lifecycle hooks: exports they replaced on already-loaded modules get their previous values back, and the mocks stop applying to modules loaded afterwards. Module mocks set up by a preload script or at the top level of a test file (including inside `describe` blocks) are the file's setup and stay in place, so `afterEach(() => mock.restore())` works alongside them. The same applies to `jest.restoreAllMocks()` and `vi.restoreAllMocks()`.
+
+A module that was loaded for the first time while a mock was active was evaluated from that mock, and `mock.restore()` cannot re-evaluate it: it keeps those exports until a later `mock.module()` replaces them again (which is what a `beforeEach` that re-installs the mock does). To load the real module in that situation, run the affected tests in a separate file.
 
 ```ts title="test.ts" icon="/icons/typescript.svg"
 import { expect, mock, spyOn, test } from "bun:test";

packages/bun-types/test.d.ts

@@ -54,6 +54,9 @@ declare module "bun:test" {
     module(id: string, factory: () => any): void | Promise<void>;
     /**
      * Restore the previous value of mocks.
+     *
+     * This also undoes `mock.module()` calls made inside tests and lifecycle hooks.
+     * Module mocks set up by a preload script or at the top level of a test file stay in place.
      */
     restore(): void;
 

src/jsc/bindings/BunPlugin.h

@@ -15,6 +15,8 @@ namespace Zig {
 
 using namespace JSC;
 
+struct ModuleMockUndoLog;
+
 class BunPlugin {
 public:
     using VirtualModuleMap = WTF::UncheckedKeyHashMap<String, JSC::Strong<JSC::JSObject>>;
@@ -79,29 +81,29 @@ class BunPlugin {
         }
 
         VirtualModuleMap* _Nullable virtualModules = nullptr;
+        // What mock.module() calls made from tests and hooks changed, so mock.restore() can undo them.
+        ModuleMockUndoLog* _Nullable moduleMockUndoLog = nullptr;
         bool mustDoExpensiveRelativeLookup = false;
         JSC::EncodedJSValue run(JSC::JSGlobalObject* globalObject, const BunString* namespaceString, const BunString* path);
 
         bool hasVirtualModules() const { return virtualModules != nullptr; }
 
         void addModuleMock(JSC::VM& vm, const String& path, JSC::JSObject* mock);
+        void restoreModuleMocks(Zig::GlobalObject* globalObject);
+        void clearVirtualModules();
+        // The log holds Strong<> roots into the realm, like the plugin lists: the test isolation swap drops it with them.
+        void discardModuleMockUndoLog();
 
         std::optional<String> resolveVirtualModule(const String& path, const String& from);
 
         void clear()
         {
             Base::clear();
-            delete virtualModules;
-            virtualModules = nullptr;
+            clearVirtualModules();
             mustDoExpensiveRelativeLookup = false;
         }
 
-        ~OnLoad()
-        {
-            if (virtualModules) {
-                delete virtualModules;
-            }
-        }
+        ~OnLoad();
     };
 
     class OnResolve final : public Base {

src/jsc/bindings/JSMockFunction.cpp

@@ -649,6 +649,15 @@ static void restoreAllMocks(Zig::GlobalObject* globalObject)
     globalObject->mockModule.activeSpies.clear();
 }
 
+JSC::JSValue unwrapSpyOriginal(JSC::JSValue value)
+{
+    if (auto* spy = dynamicDowncast<JSMockFunction>(value)) {
+        if (JSValue original = spy->spyOriginal.get())
+            return original;
+    }
+    return value;
+}
+
 extern "C" void JSMock__clearAllMocks(Zig::GlobalObject* globalObject)
 {
     // mockClear() on every mock: only clears calls, contexts, instances and results.
@@ -1480,7 +1489,10 @@ BUN_DEFINE_HOST_FUNCTION(JSMock__jsSetSystemTime, (JSC::JSGlobalObject * globalO
 BUN_DEFINE_HOST_FUNCTION(JSMock__jsRestoreAllMocks, (JSC::JSGlobalObject * globalObject, JSC::CallFrame* callframe))
 {
     auto scope = DECLARE_THROW_SCOPE(globalObject->vm());
-    restoreAllMocks(uncheckedDowncast<Zig::GlobalObject>(globalObject));
+    auto* zigGlobalObject = defaultGlobalObject(globalObject);
+    restoreAllMocks(zigGlobalObject);
+    RETURN_IF_EXCEPTION(scope, {});
+    zigGlobalObject->onLoadPlugins.restoreModuleMocks(zigGlobalObject);
     RETURN_IF_EXCEPTION(scope, {});
     return JSValue::encode(jsUndefined());
 }

src/jsc/bindings/ZigGlobalObject.cpp

@@ -710,9 +710,10 @@ extern "C" JSC::JSGlobalObject* Zig__GlobalObject__createForTestIsolation(Zig::G
         globalObject->m_processEnvObject.set(vm, globalObject, Bun::createSharedEnvironmentVariablesMap(globalObject).getObject());
     }
 
-    // The plugin registries hold Strong<> roots into the old realm; owned by the global itself,
-    // they would keep it (and everything it loaded) alive for the rest of the run.
+    // The plugin registries (and the mock.module() undo log) hold Strong<> roots into the old realm;
+    // owned by the global itself, they would keep it (and everything it loaded) alive for the rest of the run.
     oldGlobal->onLoadPlugins.clear();
+    oldGlobal->onLoadPlugins.discardModuleMockUndoLog();
     oldGlobal->onResolvePlugins.clear();
     // Drop the finished file's module registry and require.cache now rather than whenever the
     // old global happens to die. JSC's CodeCache and Bun's RuntimeTranspilerCache are VM/process

src/runtime/test_runner/jest.rs

@@ -322,6 +322,25 @@ pub(crate) mod Jest {
         unsafe { RUNNER.read() }
     }
 
+    /// `mock.restore()` keeps module mocks installed by preload or a file's collection phase (its setup), not by tests/hooks.
+    #[unsafe(no_mangle)]
+    extern "C" fn Bun__Jest__moduleMockIsPersistent(global_object: &JSGlobalObject) -> bool {
+        if global_object.bun_vm().is_in_preload {
+            return true;
+        }
+        let Some(runner) = runner_ptr() else {
+            return false;
+        };
+        // SAFETY: JS thread only; raw projections (see js_file_generation and BunTestCell::as_ptr) since we are re-entered from JS.
+        let phase = unsafe {
+            (*runner.as_ptr()).bun_test_root.active_file.as_ref().map(|file| (*file.as_ptr()).phase)
+        };
+        match phase {
+            Some(phase) => phase == bun_test::Phase::Collection,
+            None => true,
+        }
+    }
+
     #[unsafe(no_mangle)]
     extern "C" fn Bun__Jest__createTestModuleObject(
         global_object: &JSGlobalObject,

test/cli/test/isolation.test.ts

@@ -1635,6 +1635,23 @@ describe.concurrent("--isolate: collects globals pinned by leaked handles", () =
     expect(await maxLiveGlobals(String(dir))).toBeLessThanOrEqual(4);
   });
 
+  // A mock.module() made from a hook or test on a loaded module is logged so mock.restore() can undo it.
+  // The log holds the module's record, namespace and original values; left unrestored it must go with the file.
+  test("mock.module() of a loaded module from a hook, never restored", async () => {
+    using dir = tempDir("isolate-leak-module-mock-undo", {
+      ...makeLeakFixture(`
+        import { beforeAll, mock } from "bun:test";
+        import * as dep from "./dep.mjs";
+        beforeAll(() => {
+          mock.module("./dep.mjs", () => ({ getValue: () => "mocked" }));
+          if (dep.getValue() !== "mocked") throw new Error("the loaded module was not patched");
+        });
+      `),
+      "dep.mjs": `export const getValue = () => "original";`,
+    });
+    expect(await maxLiveGlobals(String(dir))).toBeLessThanOrEqual(4);
+  });
+
   test("Bun.plugin left registered", async () => {
     using dir = tempDir(
       "isolate-leak-plugin",