Merge pull request #2754 from headlamp-k8s/websocket-frontend-tests-fix

frontend: fix websocket logs errors in tests
headlamp-k8s · Jan 17, 2025 · e67e512 · e67e512
2 parents 59db42b + 39dfc3a
commit e67e512
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 37 deletions.
diff --git a/frontend/src/lib/k8s/api/v2/webSocket.test.ts b/frontend/src/lib/k8s/api/v2/webSocket.test.ts
@@ -4,7 +4,7 @@ import WS from 'vitest-websocket-mock';
 import { findKubeconfigByClusterName, getUserIdFromLocalStorage } from '../../../../stateless';
 import { getToken } from '../../../auth';
 import { getCluster } from '../../../cluster';
-import { BASE_WS_URL, useWebSocket, WebSocketManager } from './webSocket';
+import { BASE_WS_URL, MULTIPLEXER_ENDPOINT, useWebSocket, WebSocketManager } from './webSocket';
 
 // Mock dependencies
 vi.mock('../../../cluster', () => ({
@@ -37,6 +37,7 @@ describe('WebSocket Tests', () => {
   let mockServer: WS;
   let onMessage: ReturnType<typeof vi.fn>;
   let onError: ReturnType<typeof vi.fn>;
+  let originalConsoleError: typeof console.error;
 
   beforeEach(() => {
     vi.stubEnv('REACT_APP_ENABLE_WEBSOCKET_MULTIPLEXER', 'true');
@@ -48,10 +49,17 @@ describe('WebSocket Tests', () => {
     (getToken as ReturnType<typeof vi.fn>).mockReturnValue(token);
     (findKubeconfigByClusterName as ReturnType<typeof vi.fn>).mockResolvedValue({});
 
-    mockServer = new WS(`${BASE_WS_URL}wsMultiplexer`);
+    // Mock console.error for all tests
+    originalConsoleError = console.error;
+    console.error = vi.fn();
+
+    mockServer = new WS(`${BASE_WS_URL}${MULTIPLEXER_ENDPOINT}`);
   });
 
   afterEach(async () => {
+    // Restore console.error
+    console.error = originalConsoleError;
+
     WS.clean();
     vi.restoreAllMocks();
     vi.unstubAllEnvs();
@@ -211,6 +219,10 @@ describe('WebSocket Tests', () => {
       await expect(
         WebSocketManager.subscribe(clusterName, '/api/v1/pods', 'watch=true', onMessage)
       ).rejects.toThrow('WebSocket connection failed');
+
+      // Verify error was handled
+      expect(WebSocketManager.socketMultiplexer).toBeNull();
+      expect(WebSocketManager.connecting).toBe(false);
     });
 
     it('should handle duplicate subscriptions', async () => {
@@ -422,7 +434,7 @@ describe('WebSocket Tests', () => {
       expect(WebSocketManager.connecting).toBe(false);
 
       // Try to use connection again to trigger reconnect
-      const newServer = new WS(`${BASE_WS_URL}wsMultiplexer`);
+      const newServer = new WS(`${BASE_WS_URL}${MULTIPLEXER_ENDPOINT}`);
       await WebSocketManager.connect();
       await newServer.connected;
 
@@ -488,46 +500,20 @@ describe('WebSocket Tests', () => {
       consoleSpy.mockRestore();
     });
 
-    it('should handle invalid message format', async () => {
-      const path = '/api/v1/pods';
-      const query = 'watch=true';
-
-      await WebSocketManager.subscribe(clusterName, path, query, onMessage);
-      await mockServer.connected;
-      await mockServer.nextMessage; // Skip subscription
-
-      // Send invalid message
-      await mockServer.send('invalid json');
-
-      expect(onMessage).not.toHaveBeenCalled();
-    });
-
     it('should handle parse errors in message data', async () => {
-      const path = '/api/v1/pods';
-      const query = 'watch=true';
-      const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
-
-      await WebSocketManager.subscribe(clusterName, path, query, onMessage);
+      await WebSocketManager.subscribe(clusterName, '/api/v1/pods', 'watch=true', onMessage);
       await mockServer.connected;
 
-      // Skip subscription message
-      await mockServer.nextMessage;
-
-      // Send malformed data
       await mockServer.send(
         JSON.stringify({
           clusterId: clusterName,
-          path,
-          query,
-          data: 'invalid{json',
-          type: 'DATA',
+          path: '/api/v1/pods',
+          data: 'invalid json',
         })
       );
 
       expect(onMessage).not.toHaveBeenCalled();
-      expect(consoleSpy).toHaveBeenCalledWith('Failed to parse update data:', expect.any(Error));
-
-      consoleSpy.mockRestore();
+      expect(console.error).toHaveBeenCalledWith('Failed to parse update data:', expect.any(Error));
     });
 
     it('should handle message callback errors in useWebSocket', async () => {

diff --git a/frontend/src/lib/k8s/api/v2/webSocket.ts b/frontend/src/lib/k8s/api/v2/webSocket.ts
@@ -7,12 +7,12 @@ import { makeUrl } from './makeUrl';
 
 // Constants for WebSocket connection
 export const BASE_WS_URL = BASE_HTTP_URL.replace('http', 'ws');
+export const MULTIPLEXER_ENDPOINT = 'wsMultiplexer';
 
 /**
  * Multiplexer endpoint for WebSocket connections
  * This endpoint allows multiple subscriptions over a single connection
  */
-const MULTIPLEXER_ENDPOINT = 'wsMultiplexer';
 
 /**
  * Message format for WebSocket communication between client and server.
@@ -148,8 +148,8 @@ export const WebSocketManager = {
       socket.onmessage = this.handleWebSocketMessage.bind(this);
 
       socket.onerror = event => {
-        console.error('WebSocket error:', event);
         this.connecting = false;
+        console.error('WebSocket error:', event);
         reject(new Error('WebSocket connection failed'));
       };
 
@@ -297,7 +297,7 @@ export const WebSocketManager = {
     this.connecting = false;
     this.completedPaths.clear();
 
-    // Set reconnecting flag if we have active subscriptions
+    // Only log reconnecting if we have active subscriptions
     this.isReconnecting = this.activeSubscriptions.size > 0;
   },
 
@@ -318,6 +318,7 @@ export const WebSocketManager = {
       // Handle COMPLETE messages
       if (data.type === 'COMPLETE') {
         this.completedPaths.add(key);
+        return;
       }
 
       // Parse and validate update data
@@ -333,7 +334,13 @@ export const WebSocketManager = {
       if (update && typeof update === 'object') {
         const listeners = this.listeners.get(key);
         if (listeners) {
-          listeners.forEach(listener => listener(update));
+          for (const listener of listeners) {
+            try {
+              listener(update);
+            } catch (err) {
+              console.error('Failed to process WebSocket message:', err);
+            }
+          }
         }
       }
     } catch (err) {