Track current Sarus state in state machine

__tests__/index/state.test.ts

@@ -0,0 +1,68 @@
+// File Dependencies
+import Sarus from "../../src/index";
+import { WS } from "jest-websocket-mock";
+import { delay } from "../helpers/delay";
+
+const url: string = "ws://localhost:1234";
+const sarusConfig = {
+  url,
+  retryConnectionDelay: 1,
+};
+
+describe("state machine", () => {
+  it("cycles through a closed connection correctly", async () => {
+    let server: WS = new WS(url);
+
+    // In the beginning, the state is "connecting"
+    const sarus: Sarus = new Sarus(sarusConfig);
+    expect(sarus.state).toBe("connecting");
+
+    // We wait until we are connected, and see a "connected" state
+    await server.connected;
+    expect(sarus.state).toBe("connected");
+
+    // When the connection drops, the state will be "closed"
+    server.close();
+    await server.closed;
+    expect(sarus.state).toBe("closed");
+
+    // Restart server
+    server = new WS(url);
+
+    // We wait a while, and the status is "connecting" again
+    await delay(1);
+    expect(sarus.state).toBe("connecting");
+
+    // When we connect in our mock server, we are "connected" again
+    await server.connected;
+    expect(sarus.state).toBe("connected");
+
+    // Cleanup
+    server.close();
+  });
+
+  it("cycles through disconnect() correctly", async () => {
+    let server: WS = new WS(url);
+
+    // Same initial state transition as above
+    const sarus: Sarus = new Sarus(sarusConfig);
+    expect(sarus.state).toBe("connecting");
+    await server.connected;
+    expect(sarus.state).toBe("connected");
+
+    // The user can disconnect and the state will be "disconnected"
+    sarus.disconnect();
+    expect(sarus.state).toBe("disconnected");
+    await server.closed;
+
+    // The user can now reconnect, and the state will be "connecting", and then
+    // "connected" again
+    sarus.connect();
+    expect(sarus.state).toBe("connecting");
+    await server.connected;
+    // XXX for some reason the test will fail without waiting 10 ms here
+    await delay(10);
+    expect(sarus.state).toBe("connected");
+    server.close();
+  });
+});

src/index.ts

@@ -115,6 +115,29 @@ export default class Sarus {
   // Internally set
   messageStore: any;
   ws: WebSocket | undefined;
+  /*
+   * Track the current state of the Sarus object. See the diagram below.
+   *
+   *                       reconnect()    ┌──────┐
+   *      ┌───────────────────────────────│closed│
+   *      │                               └──────┘
+   *      │                                  ▲
+   *      ▼                                  │ this.ws.onclose
+   * ┌──────────┐    this.ws.onopen      ┌───┴─────┐
+   * │connecting├───────────────────────►│connected│
+   * └──────────┘                        └───┬─────┘
+   *      ▲                                  │ disconnect()
+   *      │                                  ▼
+   *      │            reconnect()      ┌────────────┐
+   *      └─────────────────────────────┤disconnected│
+   *                                    └────────────┘
+   *
+   * connect(), disconnect() are generally called by the user
+   *
+   * this.reconnect() is called internally when automatic reconnection is
+   * enabled, but can also be called by the user
+   */
+  state: "connecting" | "connected" | "disconnected" | "closed" = "connecting";
 
   constructor(props: SarusClassParams) {
     // Extract the properties that are passed to the class
@@ -293,6 +316,7 @@ export default class Sarus {
    * Connects the WebSocket client, and attaches event listeners
    */
   connect() {
+    this.state = "connecting";
     this.ws = new WebSocket(this.url, this.protocols);
     this.setBinaryType();
     this.attachEventListeners();
@@ -326,6 +350,7 @@ export default class Sarus {
    * @param {boolean} overrideDisableReconnect
    */
   disconnect(overrideDisableReconnect?: boolean) {
+    this.state = "disconnected";
     const self = this;
     // We do this to prevent automatic reconnections;
     if (!overrideDisableReconnect) {
@@ -457,7 +482,10 @@ export default class Sarus {
     WS_EVENT_NAMES.forEach((eventName) => {
       self.ws[`on${eventName}`] = (e: Function) => {
         self.eventListeners[eventName].forEach((f: Function) => f(e));
-        if (eventName === "close" && self.reconnectAutomatically) {
+        if (eventName === "open") {
+          self.state = "connected";
+        } else if (eventName === "close" && self.reconnectAutomatically) {
+          self.state = "closed";
           self.removeEventListeners();
           self.reconnect();
         }