import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
describe('WebSocket Close Event Handling', () => {
describe('AddFolderWizard WebSocket resilience', () => {
let mockWs: any;
let messageCallbacks: Array<(data: any) => void>;
let closeCallbacks: Array<() => void>;
beforeEach(() => {
messageCallbacks = [];
closeCallbacks = [];
mockWs = {
readyState: 1, // WebSocket.OPEN
send: vi.fn(),
on: vi.fn((event: string, callback: Function) => {
if (event === 'message') {
messageCallbacks.push(callback as any);
} else if (event === 'close') {
closeCallbacks.push(callback as any);
}
}),
close: vi.fn(() => {
mockWs.readyState = 3; // WebSocket.CLOSED
// Simulate close event
closeCallbacks.forEach(cb => cb());
})
};
});
afterEach(() => {
vi.clearAllMocks();
});
it('should clean up pending requests when WebSocket closes', async () => {
// Mock WebSocket client behavior similar to our updated SimpleWebSocketClient
class MockWebSocketClient {
private ws: any = null;
private messageCallbacks: Array<(data: any) => void> = [];
private closeCallbacks: Array<() => void> = [];
async connect(): Promise<void> {
this.ws = mockWs;
this.ws.on('message', (data: any) => {
this.messageCallbacks.forEach(callback => callback(data));
});
this.ws.on('close', () => {
this.ws = null;
this.closeCallbacks.forEach(callback => callback());
});
}
send(message: any): void {
if (this.isConnected()) {
this.ws.send(JSON.stringify(message));
}
}
onMessage(callback: (data: any) => void): void {
this.messageCallbacks.push(callback);
}
onClose(callback: () => void): void {
this.closeCallbacks.push(callback);
}
isConnected(): boolean {
return this.ws !== null && this.ws.readyState === 1;
}
close(): void {
if (this.ws) {
this.ws.close();
this.ws = null;
}
}
}
const client = new MockWebSocketClient();
await client.connect();
// Set up pending requests map like in AddFolderWizard
const pendingRequests = new Map<string, {
resolve: (response: any) => void;
reject: (error: Error) => void
}>();
// Set up close handler to clean pending requests
client.onClose(() => {
pendingRequests.forEach((handlers) => {
handlers.reject(new Error('WebSocket connection closed'));
});
pendingRequests.clear();
});
// Add a pending request
const testPromise = new Promise((resolve, reject) => {
pendingRequests.set('test-1', { resolve, reject });
});
expect(pendingRequests.size).toBe(1);
expect(client.isConnected()).toBe(true);
// Close the connection
client.close();
// Pending requests should be cleared
expect(pendingRequests.size).toBe(0);
expect(client.isConnected()).toBe(false);
// The promise should be rejected
await expect(testPromise).rejects.toThrow('WebSocket connection closed');
});
it('should not send messages when disconnected', () => {
class MockWebSocketClient {
private ws: any = null;
isConnected(): boolean {
return this.ws !== null && this.ws.readyState === 1;
}
send(message: any): void {
if (this.isConnected()) {
this.ws.send(JSON.stringify(message));
}
// No exception thrown, just silently doesn't send
}
setWebSocket(ws: any): void {
this.ws = ws;
}
}
const client = new MockWebSocketClient();
const mockSend = vi.fn();
// Initially disconnected
client.send({ test: 'message' });
expect(mockSend).not.toHaveBeenCalled();
// Connect
client.setWebSocket({
readyState: 1,
send: mockSend
});
client.send({ test: 'message' });
expect(mockSend).toHaveBeenCalledWith(JSON.stringify({ test: 'message' }));
// Disconnect
client.setWebSocket({
readyState: 3, // CLOSED
send: mockSend
});
mockSend.mockClear();
client.send({ test: 'another' });
expect(mockSend).not.toHaveBeenCalled();
});
it('should handle connection check before sending in request function', async () => {
const wsClient = {
isConnected: vi.fn(),
send: vi.fn(),
onMessage: vi.fn(),
onClose: vi.fn()
};
const requestModelRecommendations = async (): Promise<any[]> => {
// First check
if (!wsClient.isConnected()) {
return [];
}
return new Promise((resolve) => {
// Second check right before sending
if (!wsClient.isConnected()) {
resolve([]);
return;
}
// Would send here
wsClient.send({ type: 'test' });
resolve(['model1', 'model2']);
});
};
// Test when disconnected from start
wsClient.isConnected.mockReturnValue(false);
const result1 = await requestModelRecommendations();
expect(result1).toEqual([]);
expect(wsClient.send).not.toHaveBeenCalled();
// Test when connected
wsClient.isConnected.mockReturnValue(true);
const result2 = await requestModelRecommendations();
expect(result2).toEqual(['model1', 'model2']);
expect(wsClient.send).toHaveBeenCalledWith({ type: 'test' });
// Test when connection drops between checks
let callCount = 0;
wsClient.isConnected.mockImplementation(() => {
callCount++;
return callCount === 1; // True on first call, false on second
});
wsClient.send.mockClear();
const result3 = await requestModelRecommendations();
expect(result3).toEqual([]);
expect(wsClient.send).not.toHaveBeenCalled();
});
});
});
