refactor websocket layer

webclient/package.json

@@ -9,12 +9,14 @@
     "start": "vite",
     "preview": "vite preview",
     "test": "vitest run",
+    "test:coverage": "npm run test -- --coverage",
     "test:watch": "vitest",
     "test:integration": "vitest run --config vitest.integration.config.ts",
-    "test:integration:coverage": "vitest run --config vitest.integration.config.ts --coverage",
-    "lint": "eslint src/",
-    "lint:fix": "eslint src/ --fix",
-    "golden": "npm run lint && npm run test",
+    "test:integration:coverage": "npm run test:integration -- --coverage",
+    "lint": "eslint src",
+    "lint:fix": "eslint src --fix",
+    "golden": "npm run lint && npm run test && npm run test:integration",
+    "golden:coverage": "npm run lint && npm run test:coverage && npm run test:integration:coverage",
     "prepare": "cd .. && husky",
     "translate": "node prebuild.js -i18nOnly",
     "proto:generate": "npx buf generate"

webclient/src/websocket/WebClient.spec.ts

@@ -9,7 +9,6 @@ vi.mock('./services/WebSocketService', () => ({
     return {
       message$: { subscribe: vi.fn() },
       connect: vi.fn(),
-      testConnect: vi.fn(),
       disconnect: vi.fn(),
       send: vi.fn(),
       checkReadyState: vi.fn().mockReturnValue(true),
@@ -41,6 +40,7 @@ import { Mock } from 'vitest';
 import { SocketTransport } from './services/ProtobufService';
 import { WebSocketServiceConfig } from './services/WebSocketService';
 import type { IWebClientResponse, IWebClientRequest } from './interfaces';
+import { installMockWebSocket } from './__mocks__/helpers';
 
 function makeMockResponse(): IWebClientResponse {
   return {
@@ -90,7 +90,6 @@ describe('WebClient', () => {
       return {
         message$: messageSubject,
         connect: vi.fn(),
-        testConnect: vi.fn(),
         disconnect: vi.fn(),
         send: vi.fn(),
         checkReadyState: vi.fn().mockReturnValue(true),
@@ -156,10 +155,49 @@ describe('WebClient', () => {
   });
 
   describe('testConnect', () => {
-    it('delegates to socket.testConnect', () => {
-      const opts: Enriched.WebSocketConnectOptions = { host: 'h', port: '1', reason: App.WebSocketConnectReason.LOGIN, userName: 'u' };
+    let MockWS: ReturnType<typeof installMockWebSocket>['MockWS'];
+    let wsMockInstance: ReturnType<typeof installMockWebSocket>['mockInstance'];
+    let restoreWS: ReturnType<typeof installMockWebSocket>['restore'];
+
+    beforeEach(() => {
+      vi.useFakeTimers();
+      const installed = installMockWebSocket();
+      MockWS = installed.MockWS;
+      wsMockInstance = installed.mockInstance;
+      restoreWS = installed.restore;
+    });
+
+    afterEach(() => {
+      restoreWS();
+      vi.useRealTimers();
+    });
+
+    const opts: Enriched.WebSocketConnectOptions = { host: 'h', port: '1', reason: App.WebSocketConnectReason.LOGIN, userName: 'u' };
+
+    it('creates a WebSocket with the correct URL', () => {
       client.testConnect(opts);
-      expect(client.socket.testConnect).toHaveBeenCalledWith(opts);
+      expect(MockWS).toHaveBeenCalledWith(expect.stringContaining('://h:1'));
+    });
+
+    it('calls testConnectionSuccessful and closes on open', () => {
+      (mockResponse.session as any).testConnectionSuccessful = vi.fn();
+      client.testConnect(opts);
+      wsMockInstance.onopen();
+      expect((mockResponse.session as any).testConnectionSuccessful).toHaveBeenCalled();
+      expect(wsMockInstance.close).toHaveBeenCalled();
+    });
+
+    it('calls testConnectionFailed on error', () => {
+      (mockResponse.session as any).testConnectionFailed = vi.fn();
+      client.testConnect(opts);
+      wsMockInstance.onerror();
+      expect((mockResponse.session as any).testConnectionFailed).toHaveBeenCalled();
+    });
+
+    it('closes socket after keepalive timeout', () => {
+      client.testConnect(opts);
+      vi.advanceTimersByTime(5000);
+      expect(wsMockInstance.close).toHaveBeenCalled();
     });
   });
 

webclient/src/websocket/WebClient.ts

@@ -3,6 +3,7 @@ import { App, Enriched } from '@app/types';
 import { ProtobufService } from './services/ProtobufService';
 import { WebSocketService } from './services/WebSocketService';
 import { ping } from './commands/session';
+import { CLIENT_OPTIONS } from './config';
 import { IWebClientResponse, IWebClientRequest } from './interfaces';
 
 export class WebClient {
@@ -54,10 +55,6 @@ export class WebClient {
     WebClient._instance = this;
 
     this.response.session.initialized();
-
-    if (import.meta.env.MODE !== 'test') {
-      console.log(this);
-    }
   }
 
   public connect(options: Enriched.WebSocketConnectOptions) {
@@ -67,7 +64,24 @@ export class WebClient {
   }
 
   public testConnect(options: Enriched.WebSocketConnectOptions) {
-    this.socket.testConnect(options);
+    const protocol = window.location.hostname === 'localhost' ? 'ws' : 'wss';
+    const { host, port } = options;
+    const socket = new WebSocket(`${protocol}://${host}:${port}`);
+    socket.binaryType = 'arraybuffer';
+
+    const timeout = setTimeout(() => socket.close(), CLIENT_OPTIONS.keepalive);
+
+    socket.onopen = () => {
+      clearTimeout(timeout);
+      this.response.session.testConnectionSuccessful();
+      socket.close();
+    };
+
+    socket.onerror = () => {
+      this.response.session.testConnectionFailed();
+    };
+
+    socket.onclose = () => {};
   }
 
   public disconnect() {

webclient/src/websocket/services/KeepAliveService.spec.ts

@@ -1,22 +1,19 @@
 import { KeepAliveService } from './KeepAliveService';
-import { WebSocketService } from './WebSocketService';
 
 type KeepAliveInternal = KeepAliveService & {
   keepalivecb: NodeJS.Timeout;
   lastPingPending: boolean;
 };
 
-vi.mock('./WebSocketService');
-
 describe('KeepAliveService', () => {
   let service: KeepAliveService;
-  let mockSocket: { checkReadyState: ReturnType<typeof vi.fn> };
+  let mockIsOpen: ReturnType<typeof vi.fn>;
 
   beforeEach(() => {
     vi.useFakeTimers();
 
-    mockSocket = { checkReadyState: vi.fn().mockReturnValue(true) };
-    service = new KeepAliveService(mockSocket as unknown as WebSocketService);
+    mockIsOpen = vi.fn().mockReturnValue(true);
+    service = new KeepAliveService(mockIsOpen);
   });
 
   it('should create', () => {
@@ -28,16 +25,12 @@ describe('KeepAliveService', () => {
     let interval;
     let promise;
     let ping;
-    let checkReadyStateSpy;
 
     beforeEach(() => {
       interval = 100;
       promise = new Promise(resolve => resolvePing = resolve);
       ping = (done) => promise.then(done);
 
-      checkReadyStateSpy = vi.spyOn(mockSocket, 'checkReadyState');
-      checkReadyStateSpy.mockImplementation(() => true);
-
       service.startPingLoop(interval, ping);
       vi.advanceTimersByTime(interval);
     });
@@ -64,7 +57,7 @@ describe('KeepAliveService', () => {
 
     it('should endPingLoop if socket is not open', () => {
       vi.spyOn(service, 'endPingLoop').mockImplementation(() => {});
-      checkReadyStateSpy.mockImplementation(() => false);
+      mockIsOpen.mockReturnValue(false);
 
       resolvePing();
       vi.advanceTimersByTime(interval);

webclient/src/websocket/services/KeepAliveService.ts

@@ -1,17 +1,15 @@
 import { Subject } from 'rxjs';
 
-import { WebSocketService } from './WebSocketService';
-
 export class KeepAliveService {
-  private socket: WebSocketService;
+  private isOpen: () => boolean;
 
   private keepalivecb: NodeJS.Timeout;
   private lastPingPending: boolean;
 
   public disconnected$ = new Subject<void>();
 
-  constructor(socket: WebSocketService) {
-    this.socket = socket;
+  constructor(isOpen: () => boolean) {
+    this.isOpen = isOpen;
   }
 
   public startPingLoop(interval: number, ping: (onPong: () => void) => void): void {
@@ -22,7 +20,7 @@ export class KeepAliveService {
       }
 
       // stop the ping loop if we"re disconnected
-      if (!this.socket.checkReadyState(WebSocket.OPEN)) {
+      if (!this.isOpen()) {
         this.endPingLoop();
         return;
       }

webclient/src/websocket/services/ProtobufService.spec.ts

@@ -91,11 +91,14 @@ describe('ProtobufService', () => {
       expect(mockSocket.send).toHaveBeenCalled();
     });
 
-    it('does not send when socket is not OPEN', () => {
+    it('does not register callback or increment cmdId when transport is closed', () => {
       const service = new ProtobufService(mockSocket);
       mockSocket.isOpen.mockReturnValue(false);
-      service.sendCommand(create(Data.CommandContainerSchema), vi.fn());
+      const cb = vi.fn();
+      service.sendCommand(create(Data.CommandContainerSchema), cb);
       expect(mockSocket.send).not.toHaveBeenCalled();
+      expect((service as ProtobufInternal).cmdId).toBe(0);
+      expect((service as ProtobufInternal).pendingCommands.size).toBe(0);
     });
   });
 

webclient/src/websocket/services/ProtobufService.ts

@@ -106,14 +106,14 @@ export class ProtobufService {
   }
 
   public sendCommand(cmd: Data.CommandContainer, callback: (raw: Data.Response) => void) {
-    this.cmdId++;
+    if (!this.transport.isOpen()) {
+      return;
+    }
 
+    this.cmdId++;
     cmd.cmdId = BigInt(this.cmdId);
     this.pendingCommands.set(this.cmdId, callback);
-
-    if (this.transport.isOpen()) {
-      this.transport.send(toBinary(Data.CommandContainerSchema, cmd));
-    }
+    this.transport.send(toBinary(Data.CommandContainerSchema, cmd));
   }
 
   public handleMessageEvent({ data }: MessageEvent): void {

webclient/src/websocket/services/WebSocketService.spec.ts

@@ -13,7 +13,6 @@ import { App } from '@app/types';
 
 type WebSocketInternal = WebSocketService & {
   keepAliveService: KeepAliveService;
-  testSocket: WebSocket | null;
 };
 
 let MockWS: Mock;
@@ -23,8 +22,6 @@ let mockConfig: WebSocketServiceConfig;
 let mockResponse: {
   session: {
     connectionFailed: Mock;
-    testConnectionSuccessful: Mock;
-    testConnectionFailed: Mock;
   };
 };
 let mockOnStatusChange: Mock;
@@ -41,8 +38,6 @@ beforeEach(() => {
   mockResponse = {
     session: {
       connectionFailed: vi.fn(),
-      testConnectionSuccessful: vi.fn(),
-      testConnectionFailed: vi.fn(),
     },
   };
   mockOnStatusChange = vi.fn();
@@ -71,12 +66,6 @@ describe('WebSocketService', () => {
     return service;
   }
 
-  function createTestConnectedService() {
-    const service = new WebSocketService(mockConfig);
-    service.testConnect({ host: 'h', port: '1' }, 'ws');
-    return service;
-  }
-
   describe('constructor', () => {
     it('subscribes disconnected$ from KeepAliveService', () => {
       const service = new WebSocketService(mockConfig);
@@ -240,58 +229,4 @@ describe('WebSocketService', () => {
     });
   });
 
-  describe('testConnect', () => {
-    it('creates a test WebSocket with correct URL', () => {
-      const service = new WebSocketService(mockConfig);
-      locationRestores.push(withMockLocation({ hostname: 'example.com' }));
-      service.testConnect({ host: 'example.com', port: '9000' });
-      expect(MockWS).toHaveBeenCalledWith('wss://example.com:9000');
-    });
-
-    it('uses ws protocol on localhost', () => {
-      const service = new WebSocketService(mockConfig);
-      locationRestores.push(withMockLocation({ hostname: 'localhost' }));
-      service.testConnect({ host: 'h', port: '1' });
-      expect(MockWS).toHaveBeenCalledWith('ws://h:1');
-    });
-
-    it('closes previous testSocket when connecting again', () => {
-      const service = new WebSocketService(mockConfig);
-      service.testConnect({ host: 'h', port: '1' }, 'ws');
-      const firstInstance = mockInstance;
-      // install second mock instance and restore after test
-      const installed2 = installMockWebSocket();
-      service.testConnect({ host: 'h', port: '2' }, 'ws');
-      expect(firstInstance.close).toHaveBeenCalled();
-      // restore original mock so subsequent tests see a clean global
-      mockInstance = installed2.mockInstance;
-      MockWS = installed2.MockWS;
-    });
-
-    it('calls response.session.testConnectionSuccessful on open', () => {
-      createTestConnectedService();
-      vi.spyOn(globalThis, 'clearTimeout');
-      mockInstance.onopen();
-      expect(mockResponse.session.testConnectionSuccessful).toHaveBeenCalled();
-      expect(mockInstance.close).toHaveBeenCalled();
-    });
-
-    it('fires socket.close after keepalive timeout for testConnect', () => {
-      createTestConnectedService();
-      vi.advanceTimersByTime(1000);
-      expect(mockInstance.close).toHaveBeenCalled();
-    });
-
-    it('calls response.session.testConnectionFailed on error', () => {
-      createTestConnectedService();
-      mockInstance.onerror();
-      expect(mockResponse.session.testConnectionFailed).toHaveBeenCalled();
-    });
-
-    it('nulls out testSocket on close', () => {
-      const service = createTestConnectedService();
-      mockInstance.onclose();
-      expect((service as WebSocketInternal).testSocket).toBeNull();
-    });
-  });
 });

webclient/src/websocket/services/WebSocketService.ts

@@ -14,7 +14,6 @@ export interface WebSocketServiceConfig {
 
 export class WebSocketService {
   private socket: WebSocket;
-  private testSocket: WebSocket;
 
   private config: WebSocketServiceConfig;
   private response: IWebClientResponse;
@@ -29,7 +28,7 @@ export class WebSocketService {
     this.config = config;
     this.response = config.response;
 
-    this.keepAliveService = new KeepAliveService(this);
+    this.keepAliveService = new KeepAliveService(() => this.checkReadyState(WebSocket.OPEN));
     this.keepAliveService.disconnected$.subscribe(() => {
       this.disconnect();
       this.config.onStatusChange(App.StatusEnum.DISCONNECTED, 'Connection timeout');
@@ -47,16 +46,6 @@ export class WebSocketService {
     this.socket = this.createWebSocket(`${protocol}://${host}:${port}`);
   }
 
-  public testConnect(options: Enriched.WebSocketConnectOptions, protocol: string = 'wss'): void {
-    if (window.location.hostname === 'localhost') {
-      protocol = 'ws';
-    }
-
-    const { host, port } = options;
-
-    this.testWebSocket(`${protocol}://${host}:${port}`);
-  }
-
   public disconnect(): void {
     if (this.socket) {
       this.socket.close();
@@ -109,31 +98,4 @@ export class WebSocketService {
     return socket;
   }
 
-  private testWebSocket(url: string): void {
-    if (this.testSocket) {
-      this.testSocket.onerror = null;
-      this.testSocket.close();
-    }
-
-    const socket = new WebSocket(url);
-    socket.binaryType = 'arraybuffer';
-
-    const connectionTimer = setTimeout(() => socket.close(), CLIENT_OPTIONS.keepalive);
-
-    socket.onopen = () => {
-      clearTimeout(connectionTimer);
-      this.response.session.testConnectionSuccessful();
-      socket.close();
-    };
-
-    socket.onerror = () => {
-      this.response.session.testConnectionFailed();
-    };
-
-    socket.onclose = () => {
-      this.testSocket = null;
-    }
-
-    this.testSocket = socket;
-  }
 }