2 files changed, 96 insertions, 7 deletions
diff --git a/test/bun.js/websocket-subprocess.ts b/test/bun.js/websocket-subprocess.ts
index 4700f75b7..cf64f18ef 100644
--- a/test/bun.js/websocket-subprocess.ts
+++ b/test/bun.js/websocket-subprocess.ts
@@ -1,12 +1,13 @@
-const hostname = process.argv[2];
-const port = process.argv[3];
-
-const host = port ? `http://${hostname}:${port}` : hostname;
+const host = process.argv[2];
 
 const ws = new WebSocket(host);
 
 ws.onmessage = (message) => {
-  if (message.data == "hello websocket") {
+  if (message.data === "hello websocket") {
     ws.send("hello");
+  } else if (message.data === "timeout") {
+    setTimeout(() => {
+      ws.send("close");
+    }, 300);
   }
 };
diff --git a/test/bun.js/websocket.test.js b/test/bun.js/websocket.test.js
index 0d4a0ce06..6a30f4c11 100644
--- a/test/bun.js/websocket.test.js
+++ b/test/bun.js/websocket.test.js
@@ -158,7 +158,11 @@ describe("websocket in subprocess", () => {
       },
     });
     const subprocess = Bun.spawn({
-      cmd: [bunExe(), "websocket-subprocess.ts", server.hostname, server.port],
+      cmd: [
+        bunExe(),
+        "websocket-subprocess.ts",
+        `http://${server.hostname}:${server.port}`,
+      ],
       stderr: "pipe",
       stdin: "pipe",
       stdout: "pipe",
@@ -166,7 +170,7 @@ describe("websocket in subprocess", () => {
 
     expect(await subprocess.exited).toBe(0);
     expect(messageReceived).toBe(true);
-    server.stop();
+    server.stop(true);
   });
 
   it("should exit after killed", async () => {
@@ -181,4 +185,88 @@ describe("websocket in subprocess", () => {
 
     expect(await subprocess.exited).toBe("SIGHUP");
   });
+
+  it("should exit with invalid url", async () => {
+    const subprocess = Bun.spawn({
+      cmd: [bunExe(), "websocket-subprocess.ts", "invalid url"],
+      stderr: "pipe",
+      stdin: "pipe",
+      stdout: "pipe",
+    });
+
+    expect(await subprocess.exited).toBe(1);
+  });
+
+  it("should exit after timeout", async () => {
+    let messageReceived = false;
+    let start = 0;
+    const server = Bun.serve({
+      port: 8765,
+      fetch(req, server) {
+        if (server.upgrade(req)) {
+          return;
+        }
+
+        return new Response("http response");
+      },
+      websocket: {
+        open(ws) {
+          start = performance.now();
+          ws.send("timeout");
+        },
+        message(ws, message) {
+          messageReceived = true;
+          expect(performance.now() - start >= 300).toBe(true);
+          ws.close();
+        },
+        close(ws) {},
+      },
+    });
+    const subprocess = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "websocket-subprocess.ts",
+        `http://${server.hostname}:${server.port}`,
+      ],
+      stderr: "pipe",
+      stdin: "pipe",
+      stdout: "pipe",
+    });
+
+    expect(await subprocess.exited).toBe(0);
+    expect(messageReceived).toBe(true);
+    server.stop();
+  });
+
+  it("should exit after server stop and 0 messages", async () => {
+    const server = Bun.serve({
+      port: 8765,
+      fetch(req, server) {
+        if (server.upgrade(req)) {
+          return;
+        }
+
+        return new Response("http response");
+      },
+      websocket: {
+        open(ws) {},
+        message(ws, message) {},
+        close(ws) {},
+      },
+    });
+
+    const subprocess = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "websocket-subprocess.ts",
+        `http://${server.hostname}:${server.port}`,
+      ],
+      stderr: "pipe",
+      stdin: "pipe",
+      stdout: "pipe",
+    });
+
+    server.stop();
+    expect(await subprocess.exited).toBe(0);
+  });
 });