frontend: Add tests for websocket multiplexer

Signed-off-by: Kautilya Tripathi <[email protected]>

frontend/src/lib/k8s/api/v2/webSocket.test.ts

@@ -1,126 +1,111 @@
-import { renderHook } from '@testing-library/react';
-import { describe, expect, it, vi } from 'vitest';
-import { useWebSockets } from './webSocket';
-
-describe('useWebSockets', () => {
-  it('should return undefined when disabled', () => {
-    const { result } = renderHook(() =>
-      useWebSockets({
-        enabled: false,
-        connections: [],
-      })
-    );
-
-    expect(result.current).toBeUndefined();
-  });
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { WebSocketManager } from './webSocket';
 
-  it('should handle empty connections array', () => {
-    const { result } = renderHook(() =>
-      useWebSockets({
-        enabled: true,
-        connections: [],
-      })
-    );
+// Mock WebSocket
+class MockWebSocket {
+  onopen: (() => void) | null = null;
+  onmessage: ((event: MessageEvent) => void) | null = null;
+  onclose: (() => void) | null = null;
+  onerror: ((event: Event) => void) | null = null;
+  readyState = 0; // WebSocket.CONNECTING
+  send = vi.fn();
+  CONNECTING = 0;
+  OPEN = 1;
+  CLOSING = 2;
+  CLOSED = 3;
 
-    expect(result.current).toBeUndefined();
-  });
+  constructor() {
+    setTimeout(() => {
+      this.readyState = 1; // WebSocket.OPEN
+      this.onopen?.();
+    }, 0);
+  }
+
+  close() {
+    this.readyState = 3; // WebSocket.CLOSED
+    this.onclose?.();
+  }
+}
+
+// Mock localStorage
+const mockGetUserId = vi.fn().mockReturnValue('test-user-id');
+vi.mock('../../../helpers', () => ({
+  getUserIdFromLocalStorage: () => mockGetUserId(),
+}));
+
+describe('WebSocketManager', () => {
+  let originalWebSocket: typeof WebSocket;
+
+  beforeEach(() => {
+    // Save original WebSocket
+    originalWebSocket = global.WebSocket;
+    // Replace with mock
+    (global as any).WebSocket = MockWebSocket;
 
-  it('should handle single connection', () => {
-    const onMessage = vi.fn();
-    const { result } = renderHook(() =>
-      useWebSockets({
-        enabled: true,
-        connections: [
-          {
-            url: 'ws://localhost:3000',
-            cluster: 'test-cluster',
-            onMessage,
-          },
-        ],
-      })
-    );
-
-    expect(result.current).toBeUndefined();
+    // Reset WebSocketManager state
+    WebSocketManager.socketMultiplexer = null;
+    WebSocketManager.connecting = false;
+    WebSocketManager.isReconnecting = false;
+    WebSocketManager.listeners.clear();
+    WebSocketManager.completedPaths.clear();
+    WebSocketManager.activeSubscriptions.clear();
+    WebSocketManager.pendingUnsubscribes.clear();
+
+    // Reset mocks
+    vi.clearAllMocks();
   });
 
-  it('should handle multiple connections', () => {
-    const onMessage1 = vi.fn();
-    const onMessage2 = vi.fn();
-    const { result } = renderHook(() =>
-      useWebSockets({
-        enabled: true,
-        connections: [
-          {
-            url: 'ws://localhost:3000/path1',
-            cluster: 'test-cluster',
-            onMessage: onMessage1,
-          },
-          {
-            url: 'ws://localhost:3000/path2',
-            cluster: 'test-cluster',
-            onMessage: onMessage2,
-          },
-        ],
-      })
-    );
-
-    expect(result.current).toBeUndefined();
+  afterEach(() => {
+    // Restore original WebSocket
+    global.WebSocket = originalWebSocket;
   });
 
-  it('should handle custom protocols', () => {
-    const onMessage = vi.fn();
-    const { result } = renderHook(() =>
-      useWebSockets({
-        enabled: true,
-        protocols: ['v1.protocol', 'v2.protocol'],
-        connections: [
-          {
-            url: 'ws://localhost:3000',
-            cluster: 'test-cluster',
-            onMessage,
-          },
-        ],
-      })
-    );
-
-    expect(result.current).toBeUndefined();
+  describe('createKey', () => {
+    it('should create a unique key from clusterId, path, and query', () => {
+      const key = WebSocketManager.createKey('cluster1', '/api/v1/pods', 'watch=true');
+      expect(key).toBe('cluster1:/api/v1/pods:watch=true');
+    });
   });
 
-  it('should handle binary type', () => {
-    const onMessage = vi.fn();
-    const { result } = renderHook(() =>
-      useWebSockets({
-        enabled: true,
-        type: 'binary',
-        connections: [
-          {
-            url: 'ws://localhost:3000',
-            cluster: 'test-cluster',
-            onMessage,
-          },
-        ],
-      })
-    );
-
-    expect(result.current).toBeUndefined();
+  describe('connect', () => {
+    it('should establish a WebSocket connection', async () => {
+      const socket = await WebSocketManager.connect();
+      expect(socket).toBeDefined();
+      expect(WebSocketManager.socketMultiplexer).toBe(socket);
+      expect(WebSocketManager.connecting).toBe(false);
+    });
+
+    it('should reuse existing connection if available', async () => {
+      const socket1 = await WebSocketManager.connect();
+      const socket2 = await WebSocketManager.connect();
+      expect(socket1).toBe(socket2);
+    });
   });
 
-  it('should cleanup on unmount', () => {
-    const onMessage = vi.fn();
-    const { unmount } = renderHook(() =>
-      useWebSockets({
-        enabled: true,
-        connections: [
-          {
-            url: 'ws://localhost:3000',
-            cluster: 'test-cluster',
-            onMessage,
-          },
-        ],
-      })
-    );
-
-    unmount();
-    // If we get here without errors, cleanup worked
+  describe('handleWebSocketMessage', () => {
+    it('should handle COMPLETE messages', async () => {
+      await WebSocketManager.connect();
+      const key = WebSocketManager.createKey('cluster1', '/api/v1/pods', 'watch=true');
+
+      WebSocketManager.handleWebSocketMessage({
+        data: JSON.stringify({
+          clusterId: 'cluster1',
+          path: '/api/v1/pods',
+          query: 'watch=true',
+          type: 'COMPLETE',
+        }),
+      } as MessageEvent);
+
+      expect(WebSocketManager.completedPaths.has(key)).toBe(true);
+    });
+
+    it('should handle invalid JSON messages', async () => {
+      await WebSocketManager.connect();
+
+      // This should not throw
+      WebSocketManager.handleWebSocketMessage({
+        data: 'invalid json',
+      } as MessageEvent);
+    });
   });
 });