fix: Unset ws of WebSocketConnectorImpl after close event

neet · neet · commit 0a73a2a1fba0 · 2024-10-23T23:22:52.000+09:00
diff --git a/src/adapters/ws/web-socket-connector.spec.ts b/src/adapters/ws/web-socket-connector.spec.ts
@@ -1,6 +1,7 @@
 import getPort from "get-port";
-import { WebSocketServer } from "isomorphic-ws";
+import { WebSocket, WebSocketServer } from "isomorphic-ws";
 
+import { sleep } from "../../utils";
 import { MastoWebSocketError } from "../errors";
 import { WebSocketConnectorImpl } from "./web-socket-connector";
 
@@ -40,4 +41,35 @@ describe("WebSocketConnector", () => {
     const promise = connector.acquire();
     await expect(promise).rejects.toBeInstanceOf(MastoWebSocketError);
   });
+
+  it("creates reconnection after connection closed", async () => {
+    const port = await getPort();
+    let server = new WebSocketServer({ port });
+
+    const connector = new WebSocketConnectorImpl({
+      constructorParameters: [`ws://localhost:${port}`],
+    });
+    const connections = connector[Symbol.asyncIterator]();
+
+    // First connection
+    const { value: ws } = await connections.next();
+    await sleep(100);
+    expect(ws.readyState).toBe(WebSocket.OPEN);
+
+    // Close first connection
+    for (const client of server.clients) client.close();
+    server.close();
+    await sleep(100);
+    expect(ws.readyState).toBe(WebSocket.CLOSED);
+
+    server = new WebSocketServer({ port });
+
+    // Second connection
+    const { value: ws2 } = await connections.next();
+    await sleep(100);
+    expect(ws2.readyState).toBe(WebSocket.OPEN);
+
+    server.close();
+    connector.kill();
+  });
 });
diff --git a/src/adapters/ws/web-socket-connector.ts b/src/adapters/ws/web-socket-connector.ts
@@ -66,8 +66,6 @@ export class WebSocketConnectorImpl implements WebSocketConnector {
 
   private async spawn() {
     while (!this.killed) {
-      this.ws?.close();
-
       try {
         await this.backoff.sleep();
       } catch {
@@ -96,6 +94,8 @@ export class WebSocketConnectorImpl implements WebSocketConnector {
       } catch (error) {
         this.logger?.log("error", "WebSocket error:", error);
       }
+
+      this.ws = undefined;
     }
 
     for (const { reject } of this.queue) {
