1 files changed, 27 insertions, 7 deletions

diff --git a/test/js/web/websocket/websocket.test.js b/test/js/web/websocket/websocket.test.js
index 867b86123..76ff16ecb 100644
--- a/test/js/web/websocket/websocket.test.js
+++ b/test/js/web/websocket/websocket.test.js
@@ -6,16 +6,33 @@ const TEST_WEBSOCKET_HOST = process.env.TEST_WEBSOCKET_HOST || "wss://ws.postman
 
 describe("WebSocket", () => {
   it("should connect", async () => {
-    const ws = new WebSocket(TEST_WEBSOCKET_HOST);
-    await new Promise((resolve, reject) => {
+    const server = Bun.serve({
+      port: 0,
+      fetch(req, server) {
+        if (server.upgrade(req)) {
+          server.stop();
+          return;
+        }
+
+        return new Response();
+      },
+      websocket: {
+        open(ws) {},
+        message(ws) {
+          ws.close();
+        },
+      },
+    });
+    const ws = new WebSocket(`ws://${server.hostname}:${server.port}`, {});
+    await new Promise(resolve => {
       ws.onopen = resolve;
-      ws.onerror = reject;
     });
-    var closed = new Promise((resolve, reject) => {
+    var closed = new Promise(resolve => {
       ws.onclose = resolve;
     });
     ws.close();
     await closed;
+    server.stop(true);
   });
 
   it("should connect over https", async () => {
@@ -59,17 +76,18 @@ describe("WebSocket", () => {
     const server = Bun.serve({
       port: 0,
       fetch(req, server) {
-        server.stop();
         done();
+        server.stop();
         return new Response();
       },
       websocket: {
-        open(ws) {
+        open(ws) {},
+        message(ws) {
           ws.close();
         },
       },
     });
-    const ws = new WebSocket(`http://${server.hostname}:${server.port}`, {});
+    new WebSocket(`http://${server.hostname}:${server.port}`, {});
   });
   describe("nodebuffer", () => {
     it("should support 'nodebuffer' binaryType", done => {
@@ -93,6 +111,7 @@ describe("WebSocket", () => {
       expect(ws.binaryType).toBe("nodebuffer");
       Bun.gc(true);
       ws.onmessage = ({ data }) => {
+        ws.close();
         expect(Buffer.isBuffer(data)).toBe(true);
         expect(data).toEqual(new Uint8Array([1, 2, 3]));
         server.stop(true);
@@ -117,6 +136,7 @@ describe("WebSocket", () => {
             ws.sendBinary(new Uint8Array([1, 2, 3]));
             setTimeout(() => {
               client.onmessage = ({ data }) => {
+                client.close();
                 expect(Buffer.isBuffer(data)).toBe(true);
                 expect(data).toEqual(new Uint8Array([1, 2, 3]));
                 server.stop(true);