diff --git a/src/bun.js/bindings/BunPlugin.cpp b/src/bun.js/bindings/BunPlugin.cpp index c14c8bf33c..fccd871acc 100644 --- a/src/bun.js/bindings/BunPlugin.cpp +++ b/src/bun.js/bindings/BunPlugin.cpp @@ -407,6 +407,12 @@ public: mutable WriteBarrier callbackFunctionOrCachedResult; bool hasCalledModuleMock = false; + // Original export values snapshot (a plain JS object with properties = original exports). + // Used by mock.restore() to reverse overrideExportValue patches. + WriteBarrier originalExportsSnapshot; + // The target to restore to: JSModuleNamespaceObject (ESM) or JSCommonJSModule (CJS). + WriteBarrier restoreTarget; + static JSModuleMock* create(JSC::VM& vm, JSC::Structure* structure, JSC::JSObject* callback); static Structure* createStructure(JSC::VM& vm, JSC::JSGlobalObject* globalObject, JSC::JSValue prototype); @@ -586,6 +592,20 @@ extern "C" JSC_DEFINE_HOST_FUNCTION(JSMock__jsModuleMock, (JSC::JSGlobalObject * JSModuleMock* mock = JSModuleMock::create(vm, globalObject->mockModule.mockModuleStructure.getInitializedOnMainThread(globalObject), callback); + // If this specifier was already mocked, carry over the original snapshot + // so we preserve the true originals across re-mocks. + if (globalObject->onLoadPlugins.virtualModules) { + auto it = globalObject->onLoadPlugins.virtualModules->find(specifier); + if (it != globalObject->onLoadPlugins.virtualModules->end()) { + if (auto* existingMock = jsDynamicCast(it->value.get())) { + if (existingMock->originalExportsSnapshot) { + mock->originalExportsSnapshot.set(vm, mock, existingMock->originalExportsSnapshot.get()); + mock->restoreTarget.set(vm, mock, existingMock->restoreTarget.get()); + } + } + } + } + auto* esm = globalObject->esmRegistryMap(); RETURN_IF_EXCEPTION(scope, {}); @@ -639,6 +659,31 @@ extern "C" JSC_DEFINE_HOST_FUNCTION(JSMock__jsModuleMock, (JSC::JSGlobalObject * auto* object = exportsValue.getObject(); removeFromESM = false; + // Snapshot original export values before patching, + // so mock.restore() can reverse the overrideExportValue calls. + // We use the mock result's property names to know which exports + // to snapshot, since synthetic namespace objects may have empty m_names. + if (!mock->originalExportsSnapshot) { + mock->restoreTarget.set(vm, mock, moduleNamespaceObject); + + JSObject* snapshot = constructEmptyObject(globalObject); + if (object) { + JSC::PropertyNameArrayBuilder mockNames(vm, PropertyNameMode::Strings, PrivateSymbolMode::Exclude); + JSObject::getOwnPropertyNames(object, globalObject, mockNames, DontEnumPropertiesMode::Exclude); + RETURN_IF_EXCEPTION(scope, {}); + for (auto& name : mockNames) { + JSValue originalValue = moduleNamespaceObject->get(globalObject, name); + RETURN_IF_EXCEPTION(scope, {}); + snapshot->putDirect(vm, name, originalValue); + } + } else { + JSValue originalDefault = moduleNamespaceObject->get(globalObject, vm.propertyNames->defaultKeyword); + RETURN_IF_EXCEPTION(scope, {}); + snapshot->putDirect(vm, vm.propertyNames->defaultKeyword, originalDefault); + } + mock->originalExportsSnapshot.set(vm, mock, snapshot); + } + if (object) { JSC::PropertyNameArrayBuilder names(vm, PropertyNameMode::Strings, PrivateSymbolMode::Exclude); JSObject::getOwnPropertyNames(object, globalObject, names, DontEnumPropertiesMode::Exclude); @@ -676,6 +721,18 @@ extern "C" JSC_DEFINE_HOST_FUNCTION(JSMock__jsModuleMock, (JSC::JSGlobalObject * if (entryValue) { removeFromCJS = true; if (auto* moduleObject = entryValue ? jsDynamicCast(entryValue) : nullptr) { + // Snapshot original CJS exports before patching. + if (!mock->originalExportsSnapshot) { + JSValue currentExports = moduleObject->getIfPropertyExists(globalObject, Bun::builtinNames(vm).exportsPublicName()); + RETURN_IF_EXCEPTION(scope, {}); + if (currentExports) { + JSObject* snapshot = constructEmptyObject(globalObject); + snapshot->putDirect(vm, vm.propertyNames->defaultKeyword, currentExports); + mock->originalExportsSnapshot.set(vm, mock, snapshot); + mock->restoreTarget.set(vm, mock, moduleObject); + } + } + JSValue exportsValue = getJSValue(); RETURN_IF_EXCEPTION(scope, {}); @@ -708,10 +765,63 @@ void JSModuleMock::visitChildrenImpl(JSCell* cell, Visitor& visitor) Base::visitChildren(mock, visitor); visitor.append(mock->callbackFunctionOrCachedResult); + visitor.append(mock->originalExportsSnapshot); + visitor.append(mock->restoreTarget); } DEFINE_VISIT_CHILDREN(JSModuleMock); +void BunPlugin::OnLoad::restoreModuleMocks(JSC::JSGlobalObject* lexicalGlobalObject) +{ + if (!virtualModules) + return; + + auto& vm = lexicalGlobalObject->vm(); + auto scope = DECLARE_THROW_SCOPE(vm); + + // Collect specifiers to remove (can't mutate map while iterating). + Vector toRemove; + + for (auto& entry : *virtualModules) { + auto* moduleMock = jsDynamicCast(entry.value.get()); + if (!moduleMock) { + + continue; + } + + toRemove.append(entry.key); + + if (!moduleMock->originalExportsSnapshot || !moduleMock->restoreTarget) + continue; + + auto* snapshot = moduleMock->originalExportsSnapshot.get(); + auto* target = moduleMock->restoreTarget.get(); + + if (auto* ns = jsDynamicCast(target)) { + // ESM: restore each export value on the namespace object. + JSC::PropertyNameArrayBuilder names(vm, PropertyNameMode::Strings, PrivateSymbolMode::Exclude); + JSObject::getOwnPropertyNames(snapshot, lexicalGlobalObject, names, DontEnumPropertiesMode::Exclude); + RETURN_IF_EXCEPTION(scope, ); + + for (auto& name : names) { + JSValue originalValue = snapshot->get(lexicalGlobalObject, name); + RETURN_IF_EXCEPTION(scope, ); + ns->overrideExportValue(lexicalGlobalObject, name, originalValue); + RETURN_IF_EXCEPTION(scope, ); + } + } else if (auto* cjsModule = jsDynamicCast(target)) { + // CJS: restore original exports on the module object. + JSValue originalExports = snapshot->get(lexicalGlobalObject, vm.propertyNames->defaultKeyword); + RETURN_IF_EXCEPTION(scope, ); + cjsModule->putDirect(vm, Bun::builtinNames(vm).exportsPublicName(), originalExports, 0); + } + } + + for (auto& key : toRemove) { + virtualModules->remove(key); + } +} + EncodedJSValue BunPlugin::OnLoad::run(JSC::JSGlobalObject* globalObject, BunString* namespaceString, BunString* path) { Group* groupPtr = this->group(namespaceString ? namespaceString->toWTFString(BunString::ZeroCopy) : String()); diff --git a/src/bun.js/bindings/BunPlugin.h b/src/bun.js/bindings/BunPlugin.h index 4581383254..8e1a6f2860 100644 --- a/src/bun.js/bindings/BunPlugin.h +++ b/src/bun.js/bindings/BunPlugin.h @@ -76,6 +76,7 @@ public: bool hasVirtualModules() const { return virtualModules != nullptr; } void addModuleMock(JSC::VM& vm, const String& path, JSC::JSObject* mock); + void restoreModuleMocks(JSC::JSGlobalObject* globalObject); std::optional resolveVirtualModule(const String& path, const String& from); diff --git a/src/bun.js/bindings/JSMockFunction.cpp b/src/bun.js/bindings/JSMockFunction.cpp index 6456cff5b7..1f39b933b9 100644 --- a/src/bun.js/bindings/JSMockFunction.cpp +++ b/src/bun.js/bindings/JSMockFunction.cpp @@ -1461,7 +1461,9 @@ BUN_DEFINE_HOST_FUNCTION(JSMock__jsSetSystemTime, (JSC::JSGlobalObject * globalO BUN_DEFINE_HOST_FUNCTION(JSMock__jsRestoreAllMocks, (JSC::JSGlobalObject * globalObject, JSC::CallFrame* callframe)) { - JSMock__resetSpies(jsCast(globalObject)); + auto* zigGlobalObject = jsCast(globalObject); + JSMock__resetSpies(zigGlobalObject); + zigGlobalObject->onLoadPlugins.restoreModuleMocks(globalObject); return JSValue::encode(jsUndefined()); } diff --git a/test/js/bun/test/mock/mock-module-restore.test.ts b/test/js/bun/test/mock/mock-module-restore.test.ts new file mode 100644 index 0000000000..17ebce164a --- /dev/null +++ b/test/js/bun/test/mock/mock-module-restore.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, mock, spyOn, test } from "bun:test"; +import { fn, variable } from "./mock-module-fixture"; +import * as spyFixture from "./spymodule-fixture"; + +describe("mock.module restore", () => { + test("mock.restore() restores ESM module exports to original values", () => { + expect(fn()).toBe(42); + expect(variable).toBe(7); + + mock.module("./mock-module-fixture", () => ({ + fn: () => 999, + variable: 100, + })); + + expect(fn()).toBe(999); + expect(variable).toBe(100); + + mock.restore(); + + expect(fn()).toBe(42); + expect(variable).toBe(7); + }); + + test("re-mocking after restore works", () => { + expect(fn()).toBe(42); + expect(variable).toBe(7); + + mock.module("./mock-module-fixture", () => ({ + fn: () => 555, + variable: 55, + })); + + expect(fn()).toBe(555); + expect(variable).toBe(55); + + mock.restore(); + + expect(fn()).toBe(42); + expect(variable).toBe(7); + }); + + test("multiple re-mocks then restore goes back to true originals", () => { + expect(fn()).toBe(42); + expect(variable).toBe(7); + + mock.module("./mock-module-fixture", () => ({ + fn: () => 1, + variable: 1, + })); + expect(fn()).toBe(1); + + mock.module("./mock-module-fixture", () => ({ + fn: () => 2, + variable: 2, + })); + expect(fn()).toBe(2); + + mock.module("./mock-module-fixture", () => ({ + fn: () => 3, + variable: 3, + })); + expect(fn()).toBe(3); + + mock.restore(); + + expect(fn()).toBe(42); + expect(variable).toBe(7); + }); + + test("mock.restore() also restores spyOn alongside mock.module", () => { + const originalSpy = spyFixture.iSpy; + + spyOn(spyFixture, "iSpy"); + expect(spyFixture.iSpy).not.toBe(originalSpy); + + mock.module("./mock-module-fixture", () => ({ + fn: () => 777, + })); + expect(fn()).toBe(777); + + mock.restore(); + + expect(spyFixture.iSpy).toBe(originalSpy); + expect(fn()).toBe(42); + }); + + test("mock.restore() restores builtin modules", async () => { + const origReadFile = (await import("node:fs/promises")).readFile; + + mock.module("fs/promises", () => ({ + readFile: () => Promise.resolve("mocked-content"), + })); + + const { readFile } = await import("node:fs/promises"); + expect(await readFile("anything")).toBe("mocked-content"); + + mock.restore(); + + const { readFile: restored } = await import("node:fs/promises"); + expect(restored).toBe(origReadFile); + }); +});