From b61282e29094104c856a73fce31cf8eb591e1b5d Mon Sep 17 00:00:00 2001
From: Jarred Sumner <jarred@jarredsumner.com>
Date: Wed, 7 Feb 2024 15:40:10 -0800
Subject: [PATCH] Fix async module mocking (#8769)

Co-authored-by: Jarred Sumner <709451+Jarred-Sumner@users.noreply.github.com>
---
 src/bun.js/bindings/ModuleLoader.cpp      |  5 +++--
 src/bun.js/bindings/ModuleLoader.h        |  2 +-
 test/js/bun/test/mock/mock-module.test.ts | 10 ++++++++++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/bun.js/bindings/ModuleLoader.cpp b/src/bun.js/bindings/ModuleLoader.cpp
index 8f61ecf7b6..2cce537af3 100644
--- a/src/bun.js/bindings/ModuleLoader.cpp
+++ b/src/bun.js/bindings/ModuleLoader.cpp
@@ -169,10 +169,11 @@ void PendingVirtualModuleResult::visitChildrenImpl(JSCell* cell, Visitor& visito
 
 DEFINE_VISIT_CHILDREN(PendingVirtualModuleResult);
 
-PendingVirtualModuleResult* PendingVirtualModuleResult::create(JSC::JSGlobalObject* globalObject, const WTF::String& specifier, const WTF::String& referrer)
+PendingVirtualModuleResult* PendingVirtualModuleResult::create(JSC::JSGlobalObject* globalObject, const WTF::String& specifier, const WTF::String& referrer, bool wasModuleLock)
 {
     auto* virtualModule = create(globalObject->vm(), reinterpret_cast<Zig::GlobalObject*>(globalObject)->pendingVirtualModuleResultStructure());
     virtualModule->finishCreation(globalObject->vm(), specifier, referrer);
+    virtualModule->wasModuleMock = wasModuleLock;
     return virtualModule;
 }
 
@@ -363,7 +364,7 @@ static JSValue handleVirtualModuleResult(
         ASSERT(callData.type != CallData::Type::None);
         auto specifierString = specifier->toWTFString(BunString::ZeroCopy);
         auto referrerString = referrer->toWTFString(BunString::ZeroCopy);
-        PendingVirtualModuleResult* pendingModule = PendingVirtualModuleResult::create(globalObject, specifierString, referrerString);
+        PendingVirtualModuleResult* pendingModule = PendingVirtualModuleResult::create(globalObject, specifierString, referrerString, wasModuleMock);
         JSC::JSInternalPromise* internalPromise = pendingModule->internalPromise();
         MarkedArgumentBuffer arguments;
         arguments.append(promise);
diff --git a/src/bun.js/bindings/ModuleLoader.h b/src/bun.js/bindings/ModuleLoader.h
index 2e1bb6e0bf..8a04a2bb75 100644
--- a/src/bun.js/bindings/ModuleLoader.h
+++ b/src/bun.js/bindings/ModuleLoader.h
@@ -62,7 +62,7 @@ public:
     }
 
     JS_EXPORT_PRIVATE static PendingVirtualModuleResult* create(VM&, Structure*);
-    static PendingVirtualModuleResult* create(JSC::JSGlobalObject* globalObject, const WTF::String& specifier, const WTF::String& referrer);
+    static PendingVirtualModuleResult* create(JSC::JSGlobalObject* globalObject, const WTF::String& specifier, const WTF::String& referrer, bool wasModuleMock);
     static PendingVirtualModuleResult* createWithInitialValues(VM&, Structure*);
     static Structure* createStructure(VM&, JSGlobalObject*, JSValue);
 
diff --git a/test/js/bun/test/mock/mock-module.test.ts b/test/js/bun/test/mock/mock-module.test.ts
index 21c8d2fa5e..7d61bac2a0 100644
--- a/test/js/bun/test/mock/mock-module.test.ts
+++ b/test/js/bun/test/mock/mock-module.test.ts
@@ -11,6 +11,16 @@ import { expect, mock, spyOn, test, describe } from "bun:test";
 import { fn, iCallFn, variable, default as defaultValue, rexported, rexportedAs } from "./mock-module-fixture";
 import * as spyFixture from "./spymodule-fixture";
 
+test("mock.module async", async () => {
+  mock.module("i-am-async-and-mocked", async () => {
+    await 42;
+    await Bun.sleep(0);
+    return { a: 123 };
+  });
+
+  expect((await import("i-am-async-and-mocked")).a).toBe(123);
+});
+
 test("mock.restore", () => {
   const original = spyFixture.iSpy;
   spyOn(spyFixture, "iSpy");