Fix mock.restore() not restoring mocked modules (#7823)

This fix addresses issue #7823 where mock.restore() did not properly
restore modules that were mocked with mock.module().

Changes:
- Added JSMap properties to JSMockModule to store original ES and CJS
  module entries before mocking
- Modified JSMock__jsModuleMock to save original module entries to the
  maps before applying mocks
- Updated JSMock__jsRestoreAllMocks to restore saved modules and clear
  virtual modules (mocks)
- Added GC visitation for the new JSMaps

The fix requires re-importing modules after mock.restore() with a
cache-busting query parameter to get fresh references, as existing
import bindings cannot be mutated after creation.

Co-Authored-By: Claude <noreply@anthropic.com>

src/bun.js/bindings/BunPlugin.cpp

@@ -618,9 +618,18 @@ extern "C" JSC_DEFINE_HOST_FUNCTION(JSMock__jsModuleMock, (JSC::JSGlobalObject *
     bool removeFromESM = false;
     bool removeFromCJS = false;
 
+    // Store original modules before mocking for potential restoration
+    auto* originalESModulesMap = globalObject->mockModule.originalESModulesMap.get(globalObject);
+    auto* originalCJSModulesMap = globalObject->mockModule.originalCJSModulesMap.get(globalObject);
+
     JSValue entryValue = esm->get(globalObject, specifierString);
     RETURN_IF_EXCEPTION(scope, {});
     if (entryValue) {
+        // Save the original ES module entry before mocking
+        if (originalESModulesMap && !originalESModulesMap->has(globalObject, specifierString)) {
+            originalESModulesMap->set(globalObject, specifierString, entryValue);
+            RETURN_IF_EXCEPTION(scope, {});
+        }
         removeFromESM = true;
         JSObject* entry = entryValue ? entryValue.getObject() : nullptr;
         if (entry) {
@@ -671,6 +680,11 @@ extern "C" JSC_DEFINE_HOST_FUNCTION(JSMock__jsModuleMock, (JSC::JSGlobalObject *
     entryValue = globalObject->requireMap()->get(globalObject, specifierString);
     RETURN_IF_EXCEPTION(scope, {});
     if (entryValue) {
+        // Save the original CJS module entry before mocking
+        if (originalCJSModulesMap && !originalCJSModulesMap->has(globalObject, specifierString)) {
+            originalCJSModulesMap->set(globalObject, specifierString, entryValue);
+            RETURN_IF_EXCEPTION(scope, {});
+        }
         removeFromCJS = true;
         if (auto* moduleObject = entryValue ? jsDynamicCast<Bun::JSCommonJSModule*>(entryValue) : nullptr) {
             JSValue exportsValue = getJSValue();

src/bun.js/bindings/JSMockFunction.cpp

@@ -27,6 +27,9 @@
 #include "BunPlugin.h"
 #include "AsyncContextFrame.h"
 #include "ErrorCode.h"
+#include <JavaScriptCore/JSMap.h>
+#include <JavaScriptCore/JSMapInlines.h>
+#include <JavaScriptCore/JSMapIterator.h>
 
 BUN_DECLARE_HOST_FUNCTION(JSMock__jsUseFakeTimers);
 BUN_DECLARE_HOST_FUNCTION(JSMock__jsUseRealTimers);
@@ -777,6 +780,14 @@ JSMockModule JSMockModule::create(JSC::JSGlobalObject* globalObject)
         [](const JSC::LazyProperty<JSC::JSGlobalObject, Structure>::Initializer& init) {
             init.set(Bun::MockWithImplementationCleanupData::createStructure(init.vm, init.owner, init.owner->objectPrototype()));
         });
+    mock.originalESModulesMap.initLater(
+        [](const JSC::LazyProperty<JSC::JSGlobalObject, JSMap>::Initializer& init) {
+            init.set(JSMap::create(init.vm, init.owner->mapStructure()));
+        });
+    mock.originalCJSModulesMap.initLater(
+        [](const JSC::LazyProperty<JSC::JSGlobalObject, JSMap>::Initializer& init) {
+            init.set(JSMap::create(init.vm, init.owner->mapStructure()));
+        });
     return mock;
 }
 
@@ -1450,9 +1461,63 @@ BUN_DEFINE_HOST_FUNCTION(JSMock__jsSetSystemTime, (JSC::JSGlobalObject * globalO
     return JSValue::encode(callframe->thisValue());
 }
 
-BUN_DEFINE_HOST_FUNCTION(JSMock__jsRestoreAllMocks, (JSC::JSGlobalObject * globalObject, JSC::CallFrame* callframe))
+BUN_DEFINE_HOST_FUNCTION(JSMock__jsRestoreAllMocks, (JSC::JSGlobalObject * lexicalGlobalObject, JSC::CallFrame* callframe))
 {
-    JSMock__resetSpies(jsCast<Zig::GlobalObject*>(globalObject));
+    auto& vm = JSC::getVM(lexicalGlobalObject);
+    auto scope = DECLARE_THROW_SCOPE(vm);
+    auto* globalObject = jsCast<Zig::GlobalObject*>(lexicalGlobalObject);
+
+    JSMock__resetSpies(globalObject);
+
+    // Restore mocked modules to their original state
+    auto* originalESModulesMap = globalObject->mockModule.originalESModulesMap.get(globalObject);
+    auto* originalCJSModulesMap = globalObject->mockModule.originalCJSModulesMap.get(globalObject);
+
+    // Clear virtual modules (the mocks)
+    if (globalObject->onLoadPlugins.virtualModules) {
+        globalObject->onLoadPlugins.virtualModules->clear();
+    }
+
+    // Restore original ES modules
+    if (originalESModulesMap) {
+        auto* esm = globalObject->esmRegistryMap();
+        auto* iterator = JSMapIterator::create(globalObject, globalObject->mapIteratorStructure(), originalESModulesMap, IterationKind::Entries);
+        RETURN_IF_EXCEPTION(scope, {});
+
+        JSValue key, value;
+        while (iterator->nextKeyValue(globalObject, key, value)) {
+            RETURN_IF_EXCEPTION(scope, {});
+            if (key.isString()) {
+                esm->set(globalObject, key, value);
+                RETURN_IF_EXCEPTION(scope, {});
+            }
+        }
+
+        // Clear the map after restoring
+        originalESModulesMap->clear(globalObject);
+        RETURN_IF_EXCEPTION(scope, {});
+    }
+
+    // Restore original CJS modules
+    if (originalCJSModulesMap) {
+        auto* requireMap = globalObject->requireMap();
+        auto* iterator = JSMapIterator::create(globalObject, globalObject->mapIteratorStructure(), originalCJSModulesMap, IterationKind::Entries);
+        RETURN_IF_EXCEPTION(scope, {});
+
+        JSValue key, value;
+        while (iterator->nextKeyValue(globalObject, key, value)) {
+            RETURN_IF_EXCEPTION(scope, {});
+            if (key.isString()) {
+                requireMap->set(globalObject, key, value);
+                RETURN_IF_EXCEPTION(scope, {});
+            }
+        }
+
+        // Clear the map after restoring
+        originalCJSModulesMap->clear(globalObject);
+        RETURN_IF_EXCEPTION(scope, {});
+    }
+
     return JSValue::encode(jsUndefined());
 }
 

src/bun.js/bindings/JSMockFunction.h

@@ -28,7 +28,9 @@ public:
     V(Structure, mockModuleStructure)                \
     V(Structure, activeSpySetStructure)              \
     V(JSFunction, withImplementationCleanupFunction) \
-    V(JSC::Structure, mockWithImplementationCleanupDataStructure)
+    V(JSC::Structure, mockWithImplementationCleanupDataStructure) \
+    V(JSMap, originalESModulesMap)                   \
+    V(JSMap, originalCJSModulesMap)
 
 #define DECLARE_JSMOCKMODULE_GC_MEMBER(T, name) \
     LazyProperty<JSGlobalObject, T> name;

test/regression/issue/07823/07823-comprehensive.test.ts

@@ -0,0 +1,108 @@
+/**
+ * Tests for mock.restore() functionality with mock.module()
+ *
+ * Current implementation:
+ * - Restores ES and CJS module registry entries
+ * - Clears virtual modules (mocks)
+ * - Allows fresh imports to get original modules
+ *
+ * Limitation:
+ * - Existing module namespace objects that were mutated during mocking
+ *   are NOT automatically restored. Re-import is required.
+ */
+
+import { expect, mock, test } from "bun:test";
+
+test("mock.restore() - mock before import (ESM)", async () => {
+  // Mock before importing
+  mock.module("./esm-module", () => ({
+    esmFunction: () => "mocked",
+    esmValue: 999,
+  }));
+
+  const mocked = await import("./esm-module");
+  expect(mocked.esmFunction()).toBe("mocked");
+  expect(mocked.esmValue).toBe(999);
+
+  // Restore
+  mock.restore();
+
+  // Re-import gets original
+  const restored = await import("./esm-module?restored");
+  expect(restored.esmFunction()).toBe("original-esm");
+  expect(restored.esmValue).toBe(100);
+});
+
+test("mock.restore() - mock before require (CJS)", async () => {
+  // Mock before requiring
+  mock.module("./cjs-module", () => ({
+    cjsFunction: () => "mocked",
+  }));
+
+  const mocked = require("./cjs-module");
+  expect(mocked.cjsFunction()).toBe("mocked");
+
+  // Restore
+  mock.restore();
+
+  // Clear cache and re-require gets original
+  delete require.cache[require.resolve("./cjs-module")];
+  const restored = require("./cjs-module");
+  expect(restored.cjsFunction()).toBe("original-cjs");
+});
+
+test("mock.restore() - non-existent module", async () => {
+  // Mock a module that doesn't exist
+  mock.module("./nonexistent", () => ({
+    foo: () => "mocked",
+  }));
+
+  const mocked = await import("./nonexistent");
+  expect(mocked.foo()).toBe("mocked");
+
+  // Restore
+  mock.restore();
+
+  // Module no longer available
+  let threw = false;
+  try {
+    await import("./nonexistent");
+  } catch (e) {
+    threw = true;
+  }
+  expect(threw).toBe(true);
+});
+
+test("mock.restore() - multiple mocks", async () => {
+  // Mock multiple modules
+  mock.module("./esm-module", () => ({ esmFunction: () => "mock-esm" }));
+  mock.module("./cjs-module", () => ({ cjsFunction: () => "mock-cjs" }));
+
+  expect((await import("./esm-module")).esmFunction()).toBe("mock-esm");
+  expect(require("./cjs-module").cjsFunction()).toBe("mock-cjs");
+
+  // Restore all at once
+  mock.restore();
+
+  // Both restored
+  expect((await import("./esm-module?r2")).esmFunction()).toBe("original-esm");
+
+  delete require.cache[require.resolve("./cjs-module")];
+  expect(require("./cjs-module").cjsFunction()).toBe("original-cjs");
+});
+
+test("mock.restore() - sequential mocks", async () => {
+  // First mock
+  mock.module("./esm-module", () => ({ esmFunction: () => "mock1" }));
+  expect((await import("./esm-module")).esmFunction()).toBe("mock1");
+
+  mock.restore();
+  expect((await import("./esm-module?seq1")).esmFunction()).toBe("original-esm");
+
+  // Second mock
+  mock.module("./esm-module", () => ({ esmFunction: () => "mock2" }));
+  expect((await import("./esm-module")).esmFunction()).toBe("mock2");
+
+  mock.restore();
+  expect((await import("./esm-module?seq2")).esmFunction()).toBe("original-esm");
+});

test/regression/issue/07823/07823.fixture.ts

@@ -0,0 +1,3 @@
+export function foo() {
+  return "original";
+}

test/regression/issue/07823/07823.test.ts

@@ -0,0 +1,22 @@
+import { expect, mock, test } from "bun:test";
+import { foo } from "./07823.fixture";
+
+test("mock.restore() should restore original module behavior", async () => {
+  // before mock - should return "original"
+  expect(foo()).toBe("original");
+
+  // mock the module
+  mock.module("./07823.fixture", () => ({
+    foo: () => "mocked",
+  }));
+
+  // after mock - should return "mocked"
+  expect(foo()).toBe("mocked");
+
+  // restore mock
+  mock.restore();
+
+  // after restore, reimport to get fresh module - should return "original" again
+  const restored = await import("./07823.fixture?t=" + Date.now());
+  expect(restored.foo()).toBe("original");
+});

test/regression/issue/07823/cjs-module.js

@@ -0,0 +1,4 @@
+module.exports = {
+  cjsFunction: () => "original-cjs",
+  cjsValue: 42,
+};

test/regression/issue/07823/esm-module.ts

@@ -0,0 +1,7 @@
+export function esmFunction() {
+  return "original-esm";
+}
+
+export const esmValue = 100;
+
+export default "default-esm";