Add WebSocket reconnection tests (68 unit + 18 integration)
This commit is contained in:
@@ -5,7 +5,7 @@
|
||||
* Tests modal open/close, configuration editing, saving, and backup/restore
|
||||
*/
|
||||
|
||||
import { test, expect } from '@playwright/test';
|
||||
import { expect, test } from '@playwright/test';
|
||||
|
||||
test.describe('Settings Modal - Basic Functionality', () => {
|
||||
test.beforeEach(async ({ page }) => {
|
||||
|
||||
922
tests/frontend/unit/websocket.test.js
Normal file
922
tests/frontend/unit/websocket.test.js
Normal file
@@ -0,0 +1,922 @@
|
||||
/**
|
||||
* Unit tests for WebSocket client functionality
|
||||
* Tests connection, reconnection, authentication, error handling, and message dispatch
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
|
||||
|
||||
// Mock WebSocket class
|
||||
class MockWebSocket {
|
||||
constructor(url) {
|
||||
this.url = url;
|
||||
this.readyState = MockWebSocket.CONNECTING;
|
||||
this.CONNECTING = 0;
|
||||
this.OPEN = 1;
|
||||
this.CLOSING = 2;
|
||||
this.CLOSED = 3;
|
||||
|
||||
// Event handlers
|
||||
this.onopen = null;
|
||||
this.onclose = null;
|
||||
this.onerror = null;
|
||||
this.onmessage = null;
|
||||
|
||||
// Store instance for testing
|
||||
MockWebSocket._lastInstance = this;
|
||||
|
||||
// Auto-connect after a tick
|
||||
setTimeout(() => {
|
||||
if (this.readyState === MockWebSocket.CONNECTING) {
|
||||
this.readyState = MockWebSocket.OPEN;
|
||||
if (this.onopen) this.onopen({ type: 'open' });
|
||||
}
|
||||
}, 0);
|
||||
}
|
||||
|
||||
send(data) {
|
||||
if (this.readyState !== MockWebSocket.OPEN) {
|
||||
throw new Error('WebSocket is not open');
|
||||
}
|
||||
this._lastSent = data;
|
||||
}
|
||||
|
||||
close(code, reason) {
|
||||
this.readyState = MockWebSocket.CLOSING;
|
||||
setTimeout(() => {
|
||||
this.readyState = MockWebSocket.CLOSED;
|
||||
if (this.onclose) {
|
||||
this.onclose({
|
||||
type: 'close',
|
||||
code: code || 1000,
|
||||
reason: reason || '',
|
||||
wasClean: code === 1000
|
||||
});
|
||||
}
|
||||
}, 0);
|
||||
}
|
||||
|
||||
// Test helper: simulate message received
|
||||
_simulateMessage(data) {
|
||||
if (this.onmessage) {
|
||||
this.onmessage({
|
||||
type: 'message',
|
||||
data: typeof data === 'string' ? data : JSON.stringify(data)
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Test helper: simulate error
|
||||
_simulateError(error) {
|
||||
if (this.onerror) {
|
||||
this.onerror({
|
||||
type: 'error',
|
||||
error: error || new Error('WebSocket error')
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Test helper: simulate connection close
|
||||
_simulateClose(code = 1006, reason = 'Connection lost') {
|
||||
this.readyState = MockWebSocket.CLOSED;
|
||||
if (this.onclose) {
|
||||
this.onclose({
|
||||
type: 'close',
|
||||
code,
|
||||
reason,
|
||||
wasClean: false
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Static properties
|
||||
MockWebSocket.CONNECTING = 0;
|
||||
MockWebSocket.OPEN = 1;
|
||||
MockWebSocket.CLOSING = 2;
|
||||
MockWebSocket.CLOSED = 3;
|
||||
|
||||
// Import WebSocket client code (we'll need to evaluate it)
|
||||
// For testing, we'll load the actual file
|
||||
let WebSocketClient;
|
||||
|
||||
describe('WebSocket Client - Initialization', () => {
|
||||
beforeEach(() => {
|
||||
// Mock global WebSocket
|
||||
global.WebSocket = MockWebSocket;
|
||||
|
||||
// Clear any timers
|
||||
vi.useFakeTimers();
|
||||
|
||||
// Load WebSocketClient class by evaluating the source
|
||||
// In a real setup, this would be imported
|
||||
const sourceCode = `
|
||||
class WebSocketClient {
|
||||
constructor(url, options = {}) {
|
||||
this.url = url;
|
||||
this.ws = null;
|
||||
this.isConnected = false;
|
||||
this.reconnectAttempts = 0;
|
||||
this.maxReconnectAttempts = options.maxReconnectAttempts || 5;
|
||||
this.reconnectDelay = options.reconnectDelay || 1000;
|
||||
this.autoReconnect = options.autoReconnect !== false;
|
||||
this.eventHandlers = new Map();
|
||||
this.messageQueue = [];
|
||||
this.rooms = new Set();
|
||||
}
|
||||
|
||||
getWebSocketUrl() {
|
||||
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
const host = window.location.host;
|
||||
return \`\${protocol}//\${host}\${this.url}\`;
|
||||
}
|
||||
|
||||
connect() {
|
||||
try {
|
||||
const wsUrl = this.getWebSocketUrl();
|
||||
this.ws = new WebSocket(wsUrl);
|
||||
|
||||
this.ws.onopen = (event) => {
|
||||
this.isConnected = true;
|
||||
this.reconnectAttempts = 0;
|
||||
this.emit('connect');
|
||||
this.rejoinRooms();
|
||||
this.processMessageQueue();
|
||||
};
|
||||
|
||||
this.ws.onmessage = (event) => {
|
||||
this.handleMessage(event);
|
||||
};
|
||||
|
||||
this.ws.onerror = (event) => {
|
||||
console.error('WebSocket error:', event);
|
||||
this.emit('error', event.error || new Error('WebSocket error'));
|
||||
};
|
||||
|
||||
this.ws.onclose = (event) => {
|
||||
this.isConnected = false;
|
||||
this.emit('disconnect', event.reason);
|
||||
|
||||
if (this.autoReconnect && !event.wasClean &&
|
||||
this.reconnectAttempts < this.maxReconnectAttempts) {
|
||||
this.reconnectAttempts++;
|
||||
const delay = this.reconnectDelay * this.reconnectAttempts;
|
||||
console.log(\`Reconnecting in \${delay}ms (attempt \${this.reconnectAttempts}/\${this.maxReconnectAttempts})...\`);
|
||||
setTimeout(() => this.connect(), delay);
|
||||
} else if (this.reconnectAttempts >= this.maxReconnectAttempts) {
|
||||
this.emit('reconnect_failed');
|
||||
}
|
||||
};
|
||||
} catch (error) {
|
||||
console.error('Failed to create WebSocket:', error);
|
||||
this.emit('error', error);
|
||||
}
|
||||
}
|
||||
|
||||
disconnect() {
|
||||
if (this.ws) {
|
||||
this.autoReconnect = false;
|
||||
this.ws.close(1000, 'Client disconnect');
|
||||
}
|
||||
}
|
||||
|
||||
handleMessage(event) {
|
||||
try {
|
||||
const message = JSON.parse(event.data);
|
||||
const { type, ...data } = message;
|
||||
if (type) {
|
||||
this.emit(type, data);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Failed to parse message:', error);
|
||||
this.emit('error', error);
|
||||
}
|
||||
}
|
||||
|
||||
on(event, handler) {
|
||||
if (!this.eventHandlers.has(event)) {
|
||||
this.eventHandlers.set(event, []);
|
||||
}
|
||||
this.eventHandlers.get(event).push(handler);
|
||||
}
|
||||
|
||||
off(event, handler) {
|
||||
if (this.eventHandlers.has(event)) {
|
||||
const handlers = this.eventHandlers.get(event);
|
||||
const index = handlers.indexOf(handler);
|
||||
if (index !== -1) {
|
||||
handlers.splice(index, 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
emit(event, data) {
|
||||
if (this.eventHandlers.has(event)) {
|
||||
this.eventHandlers.get(event).forEach(handler => {
|
||||
try {
|
||||
handler(data);
|
||||
} catch (error) {
|
||||
console.error(\`Error in event handler for '\${event}':\`, error);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
send(action, data) {
|
||||
const message = JSON.stringify({ action, ...data });
|
||||
if (this.connected()) {
|
||||
this.ws.send(message);
|
||||
} else {
|
||||
this.messageQueue.push(message);
|
||||
}
|
||||
}
|
||||
|
||||
join(room) {
|
||||
this.rooms.add(room);
|
||||
if (this.connected()) {
|
||||
this.send('join', { room });
|
||||
}
|
||||
}
|
||||
|
||||
leave(room) {
|
||||
this.rooms.delete(room);
|
||||
if (this.connected()) {
|
||||
this.send('leave', { room });
|
||||
}
|
||||
}
|
||||
|
||||
rejoinRooms() {
|
||||
this.rooms.forEach(room => {
|
||||
this.send('join', { room });
|
||||
});
|
||||
}
|
||||
|
||||
processMessageQueue() {
|
||||
while (this.messageQueue.length > 0 && this.connected()) {
|
||||
const message = this.messageQueue.shift();
|
||||
this.ws.send(message);
|
||||
}
|
||||
}
|
||||
|
||||
connected() {
|
||||
return this.isConnected && this.ws && this.ws.readyState === WebSocket.OPEN;
|
||||
}
|
||||
}
|
||||
|
||||
function io(url) {
|
||||
const client = new WebSocketClient(url);
|
||||
client.connect();
|
||||
return client;
|
||||
}
|
||||
|
||||
globalThis.WebSocketClient = WebSocketClient;
|
||||
globalThis.io = io;
|
||||
`;
|
||||
|
||||
eval(sourceCode);
|
||||
WebSocketClient = globalThis.WebSocketClient;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
delete globalThis.WebSocketClient;
|
||||
delete globalThis.io;
|
||||
});
|
||||
|
||||
it('should create WebSocket client with default options', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
expect(client.url).toBe('/ws');
|
||||
expect(client.maxReconnectAttempts).toBe(5);
|
||||
expect(client.reconnectDelay).toBe(1000);
|
||||
expect(client.autoReconnect).toBe(true);
|
||||
expect(client.isConnected).toBe(false);
|
||||
expect(client.reconnectAttempts).toBe(0);
|
||||
});
|
||||
|
||||
it('should create WebSocket client with custom options', () => {
|
||||
const client = new WebSocketClient('/ws', {
|
||||
maxReconnectAttempts: 10,
|
||||
reconnectDelay: 2000,
|
||||
autoReconnect: false
|
||||
});
|
||||
|
||||
expect(client.maxReconnectAttempts).toBe(10);
|
||||
expect(client.reconnectDelay).toBe(2000);
|
||||
expect(client.autoReconnect).toBe(false);
|
||||
});
|
||||
|
||||
it('should initialize empty event handlers map', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
expect(client.eventHandlers).toBeInstanceOf(Map);
|
||||
expect(client.eventHandlers.size).toBe(0);
|
||||
});
|
||||
|
||||
it('should initialize empty message queue', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
expect(client.messageQueue).toEqual([]);
|
||||
});
|
||||
|
||||
it('should initialize empty rooms set', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
expect(client.rooms).toBeInstanceOf(Set);
|
||||
expect(client.rooms.size).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Connection', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
|
||||
// Mock window.location
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
|
||||
const sourceCode = `${/* Same source as above */}`;
|
||||
eval(sourceCode);
|
||||
WebSocketClient = globalThis.WebSocketClient;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
delete global.window;
|
||||
});
|
||||
|
||||
it('should generate correct WebSocket URL with http protocol', () => {
|
||||
const client = new WebSocketClient('/ws/updates');
|
||||
expect(client.getWebSocketUrl()).toBe('ws://localhost:8000/ws/updates');
|
||||
});
|
||||
|
||||
it('should generate correct WebSocket URL with https protocol', () => {
|
||||
global.window.location.protocol = 'https:';
|
||||
const client = new WebSocketClient('/ws/updates');
|
||||
expect(client.getWebSocketUrl()).toBe('wss://localhost:8000/ws/updates');
|
||||
});
|
||||
|
||||
it('should create WebSocket connection on connect()', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
|
||||
expect(MockWebSocket._lastInstance).toBeDefined();
|
||||
expect(MockWebSocket._lastInstance.url).toBe('ws://localhost:8000/ws');
|
||||
});
|
||||
|
||||
it('should emit connect event when connection opens', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const connectHandler = vi.fn();
|
||||
client.on('connect', connectHandler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(connectHandler).toHaveBeenCalledTimes(1);
|
||||
expect(client.isConnected).toBe(true);
|
||||
});
|
||||
|
||||
it('should reset reconnect attempts on successful connection', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.reconnectAttempts = 3;
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.reconnectAttempts).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Reconnection Logic', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
delete global.window;
|
||||
});
|
||||
|
||||
it('should attempt reconnection after unclean close', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Simulate connection loss
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006, 'Connection lost');
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Should trigger reconnection after delay
|
||||
expect(client.reconnectAttempts).toBe(1);
|
||||
});
|
||||
|
||||
it('should use exponential backoff for reconnection delays', async () => {
|
||||
const client = new WebSocketClient('/ws', { reconnectDelay: 1000 });
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// First reconnection attempt: 1000ms delay
|
||||
let ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006);
|
||||
|
||||
vi.advanceTimersByTime(999);
|
||||
expect(client.reconnectAttempts).toBe(1);
|
||||
|
||||
vi.advanceTimersByTime(1);
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Second reconnection attempt: 2000ms delay
|
||||
ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006);
|
||||
|
||||
vi.advanceTimersByTime(1999);
|
||||
expect(client.reconnectAttempts).toBe(2);
|
||||
|
||||
vi.advanceTimersByTime(1);
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Third reconnection attempt: 3000ms delay
|
||||
expect(client.reconnectAttempts).toBe(3);
|
||||
});
|
||||
|
||||
it('should stop reconnecting after max attempts', async () => {
|
||||
const client = new WebSocketClient('/ws', {
|
||||
maxReconnectAttempts: 3,
|
||||
reconnectDelay: 100
|
||||
});
|
||||
const reconnectFailedHandler = vi.fn();
|
||||
client.on('reconnect_failed', reconnectFailedHandler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Simulate 3 connection failures
|
||||
for (let i = 0; i < 3; i++) {
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006);
|
||||
await vi.runAllTimersAsync();
|
||||
}
|
||||
|
||||
expect(client.reconnectAttempts).toBe(3);
|
||||
expect(reconnectFailedHandler).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Should not attempt another reconnection
|
||||
const attemptsBefore = client.reconnectAttempts;
|
||||
await vi.advanceTimersByTimeAsync(5000);
|
||||
expect(client.reconnectAttempts).toBe(attemptsBefore);
|
||||
});
|
||||
|
||||
it('should not reconnect after clean disconnect', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const attemptsBefore = client.reconnectAttempts;
|
||||
client.disconnect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.reconnectAttempts).toBe(attemptsBefore);
|
||||
expect(client.autoReconnect).toBe(false);
|
||||
});
|
||||
|
||||
it('should not reconnect when autoReconnect is disabled', async () => {
|
||||
const client = new WebSocketClient('/ws', { autoReconnect: false });
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006);
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.reconnectAttempts).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Event Handling', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should register event handlers', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const handler = vi.fn();
|
||||
|
||||
client.on('test_event', handler);
|
||||
|
||||
expect(client.eventHandlers.has('test_event')).toBe(true);
|
||||
expect(client.eventHandlers.get('test_event')).toContain(handler);
|
||||
});
|
||||
|
||||
it('should register multiple handlers for same event', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const handler1 = vi.fn();
|
||||
const handler2 = vi.fn();
|
||||
|
||||
client.on('test_event', handler1);
|
||||
client.on('test_event', handler2);
|
||||
|
||||
expect(client.eventHandlers.get('test_event').length).toBe(2);
|
||||
});
|
||||
|
||||
it('should emit events to registered handlers', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const handler = vi.fn();
|
||||
|
||||
client.on('test_event', handler);
|
||||
client.emit('test_event', { message: 'test' });
|
||||
|
||||
expect(handler).toHaveBeenCalledWith({ message: 'test' });
|
||||
});
|
||||
|
||||
it('should remove event handlers with off()', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const handler = vi.fn();
|
||||
|
||||
client.on('test_event', handler);
|
||||
client.off('test_event', handler);
|
||||
|
||||
client.emit('test_event', { message: 'test' });
|
||||
expect(handler).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle errors in event handlers gracefully', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const errorHandler = vi.fn(() => {
|
||||
throw new Error('Handler error');
|
||||
});
|
||||
const normalHandler = vi.fn();
|
||||
|
||||
client.on('test_event', errorHandler);
|
||||
client.on('test_event', normalHandler);
|
||||
|
||||
// Should not throw, should continue to next handler
|
||||
expect(() => client.emit('test_event', {})).not.toThrow();
|
||||
expect(errorHandler).toHaveBeenCalled();
|
||||
expect(normalHandler).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Message Handling', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should parse and emit JSON messages', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const handler = vi.fn();
|
||||
client.on('download_progress', handler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateMessage({
|
||||
type: 'download_progress',
|
||||
episode_id: '123',
|
||||
progress: 50
|
||||
});
|
||||
|
||||
expect(handler).toHaveBeenCalledWith({
|
||||
episode_id: '123',
|
||||
progress: 50
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle malformed JSON messages', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const errorHandler = vi.fn();
|
||||
client.on('error', errorHandler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateMessage('not valid json{');
|
||||
|
||||
expect(errorHandler).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should emit error for messages without type', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const testHandler = vi.fn();
|
||||
client.on('test', testHandler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateMessage({ data: 'no type field' });
|
||||
|
||||
// Should not emit to any handler without type
|
||||
expect(testHandler).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Message Queueing', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should queue messages when disconnected', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
client.send('test_action', { data: 'test' });
|
||||
|
||||
expect(client.messageQueue.length).toBe(1);
|
||||
expect(client.messageQueue[0]).toContain('test_action');
|
||||
});
|
||||
|
||||
it('should send messages immediately when connected', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
client.send('test_action', { data: 'test' });
|
||||
|
||||
expect(ws._lastSent).toContain('test_action');
|
||||
expect(client.messageQueue.length).toBe(0);
|
||||
});
|
||||
|
||||
it('should process queued messages on reconnection', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
// Queue messages while disconnected
|
||||
client.send('action1', { data: '1' });
|
||||
client.send('action2', { data: '2' });
|
||||
|
||||
expect(client.messageQueue.length).toBe(2);
|
||||
|
||||
// Connect and process queue
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.messageQueue.length).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Room Management', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should add room to rooms set on join', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
|
||||
client.join('downloads');
|
||||
|
||||
expect(client.rooms.has('downloads')).toBe(true);
|
||||
});
|
||||
|
||||
it('should send join message when connected', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
client.join('downloads');
|
||||
|
||||
expect(ws._lastSent).toContain('join');
|
||||
expect(ws._lastSent).toContain('downloads');
|
||||
});
|
||||
|
||||
it('should remove room from rooms set on leave', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.join('downloads');
|
||||
|
||||
client.leave('downloads');
|
||||
|
||||
expect(client.rooms.has('downloads')).toBe(false);
|
||||
});
|
||||
|
||||
it('should rejoin all rooms on reconnection', async () => {
|
||||
const client = new WebSocketClient('/ws', { reconnectDelay: 100 });
|
||||
client.join('downloads');
|
||||
client.join('progress');
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Simulate disconnect and reconnect
|
||||
let ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006);
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
// Check that join messages were sent for both rooms
|
||||
ws = MockWebSocket._lastInstance;
|
||||
const sentMessages = [];
|
||||
const originalSend = ws.send.bind(ws);
|
||||
ws.send = (data) => {
|
||||
sentMessages.push(data);
|
||||
return originalSend(data);
|
||||
};
|
||||
|
||||
// Trigger rejoin
|
||||
client.rejoinRooms();
|
||||
|
||||
expect(sentMessages.some(msg => msg.includes('downloads'))).toBe(true);
|
||||
expect(sentMessages.some(msg => msg.includes('progress'))).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Error Handling', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should emit error event on WebSocket error', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const errorHandler = vi.fn();
|
||||
client.on('error', errorHandler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateError(new Error('Connection failed'));
|
||||
|
||||
expect(errorHandler).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should emit disconnect event on connection close', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
const disconnectHandler = vi.fn();
|
||||
client.on('disconnect', disconnectHandler);
|
||||
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1006, 'Connection lost');
|
||||
|
||||
expect(disconnectHandler).toHaveBeenCalledWith('Connection lost');
|
||||
});
|
||||
|
||||
it('should set isConnected to false on close', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.isConnected).toBe(true);
|
||||
|
||||
const ws = MockWebSocket._lastInstance;
|
||||
ws._simulateClose(1000);
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.isConnected).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Connection State', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should return false when not connected', () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
expect(client.connected()).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true when connected', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.connected()).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false after disconnection', async () => {
|
||||
const client = new WebSocketClient('/ws');
|
||||
client.connect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
client.disconnect();
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client.connected()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WebSocket Client - Socket.IO Compatibility', () => {
|
||||
beforeEach(() => {
|
||||
global.WebSocket = MockWebSocket;
|
||||
vi.useFakeTimers();
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: 'localhost:8000'
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('should create and connect client using io() function', async () => {
|
||||
const client = globalThis.io('/ws');
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
expect(client).toBeInstanceOf(WebSocketClient);
|
||||
expect(client.isConnected).toBe(true);
|
||||
});
|
||||
|
||||
it('should support Socket.IO-like event interface', async () => {
|
||||
const client = globalThis.io('/ws');
|
||||
const handler = vi.fn();
|
||||
client.on('test', handler);
|
||||
|
||||
await vi.runAllTimersAsync();
|
||||
|
||||
client.emit('test', { data: 'value' });
|
||||
expect(handler).toHaveBeenCalledWith({ data: 'value' });
|
||||
});
|
||||
});
|
||||
@@ -504,7 +504,7 @@ class TestBackupEdgeCases:
|
||||
async def test_concurrent_backup_operations(self, authenticated_client):
|
||||
"""Test multiple concurrent backup operations."""
|
||||
import asyncio
|
||||
|
||||
|
||||
# Create multiple backups concurrently
|
||||
tasks = [
|
||||
authenticated_client.post("/api/config/backups")
|
||||
|
||||
658
tests/integration/test_websocket_resilience.py
Normal file
658
tests/integration/test_websocket_resilience.py
Normal file
@@ -0,0 +1,658 @@
|
||||
"""Integration tests for WebSocket resilience and stress testing.
|
||||
|
||||
This module tests WebSocket connection resilience, concurrent client handling,
|
||||
server restart recovery, authentication, message ordering, and broadcast filtering.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import List, Dict, Any
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import WebSocket
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from src.server.services.websocket_service import WebSocketService, get_websocket_service
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def websocket_service():
|
||||
"""Create a WebSocketService instance for testing."""
|
||||
return WebSocketService()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_token():
|
||||
"""Create a mock authentication token for testing."""
|
||||
return "test_auth_token_12345"
|
||||
|
||||
|
||||
class MockWebSocketClient:
|
||||
"""Mock WebSocket client for testing."""
|
||||
|
||||
def __init__(self, client_id: str, service: WebSocketService):
|
||||
self.client_id = client_id
|
||||
self.service = service
|
||||
self.received_messages: List[Dict[str, Any]] = []
|
||||
self.is_connected = False
|
||||
self.websocket = Mock(spec=WebSocket)
|
||||
self.websocket.send_json = self._mock_send_json
|
||||
self.websocket.accept = self._mock_accept
|
||||
|
||||
async def _mock_accept(self):
|
||||
"""Mock WebSocket accept."""
|
||||
self.is_connected = True
|
||||
|
||||
async def _mock_send_json(self, data: Dict[str, Any]):
|
||||
"""Mock WebSocket send_json to capture messages."""
|
||||
self.received_messages.append(data)
|
||||
|
||||
async def connect(self, metadata: Dict[str, Any] = None):
|
||||
"""Connect the mock client to the service."""
|
||||
await self.service._manager.connect(
|
||||
self.websocket,
|
||||
self.client_id,
|
||||
metadata or {}
|
||||
)
|
||||
self.is_connected = True
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect the mock client from the service."""
|
||||
await self.service._manager.disconnect(self.client_id)
|
||||
self.is_connected = False
|
||||
|
||||
async def join_room(self, room: str):
|
||||
"""Join a room."""
|
||||
await self.service._manager.join_room(self.client_id, room)
|
||||
|
||||
async def leave_room(self, room: str):
|
||||
"""Leave a room."""
|
||||
await self.service._manager.leave_room(self.client_id, room)
|
||||
|
||||
def clear_messages(self):
|
||||
"""Clear received messages."""
|
||||
self.received_messages.clear()
|
||||
|
||||
|
||||
class TestWebSocketConcurrentClients:
|
||||
"""Test WebSocket handling of multiple concurrent clients."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_concurrent_connections(self, websocket_service):
|
||||
"""Test handling 100+ concurrent WebSocket clients."""
|
||||
num_clients = 100
|
||||
clients: List[MockWebSocketClient] = []
|
||||
|
||||
# Connect 100 clients
|
||||
for i in range(num_clients):
|
||||
client = MockWebSocketClient(f"client_{i}", websocket_service)
|
||||
await client.connect({"user_id": f"user_{i}"})
|
||||
clients.append(client)
|
||||
|
||||
# Verify all clients are connected
|
||||
assert len(websocket_service._manager._active_connections) == num_clients
|
||||
|
||||
# Broadcast a message to all clients
|
||||
test_message = {
|
||||
"type": "test_broadcast",
|
||||
"timestamp": time.time(),
|
||||
"message": "Test broadcast to all clients"
|
||||
}
|
||||
await websocket_service.broadcast(test_message)
|
||||
|
||||
# Verify all clients received the message
|
||||
for client in clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == test_message
|
||||
|
||||
# Disconnect all clients
|
||||
for client in clients:
|
||||
await client.disconnect()
|
||||
|
||||
assert len(websocket_service._manager._active_connections) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_room_broadcasts(self, websocket_service):
|
||||
"""Test broadcasting to specific rooms with concurrent clients."""
|
||||
# Create clients in different rooms
|
||||
room_a_clients = []
|
||||
room_b_clients = []
|
||||
room_both_clients = []
|
||||
|
||||
for i in range(10):
|
||||
client = MockWebSocketClient(f"room_a_{i}", websocket_service)
|
||||
await client.connect()
|
||||
await client.join_room("room_a")
|
||||
room_a_clients.append(client)
|
||||
|
||||
for i in range(10):
|
||||
client = MockWebSocketClient(f"room_b_{i}", websocket_service)
|
||||
await client.connect()
|
||||
await client.join_room("room_b")
|
||||
room_b_clients.append(client)
|
||||
|
||||
for i in range(5):
|
||||
client = MockWebSocketClient(f"room_both_{i}", websocket_service)
|
||||
await client.connect()
|
||||
await client.join_room("room_a")
|
||||
await client.join_room("room_b")
|
||||
room_both_clients.append(client)
|
||||
|
||||
# Broadcast to room_a
|
||||
message_a = {"type": "room_a_message", "data": "Message for room A"}
|
||||
await websocket_service._manager.broadcast_to_room(message_a, "room_a")
|
||||
|
||||
# Verify room_a and room_both clients received, room_b did not
|
||||
for client in room_a_clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == message_a
|
||||
|
||||
for client in room_both_clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == message_a
|
||||
|
||||
for client in room_b_clients:
|
||||
assert len(client.received_messages) == 0
|
||||
|
||||
# Clear messages
|
||||
for client in room_a_clients + room_b_clients + room_both_clients:
|
||||
client.clear_messages()
|
||||
|
||||
# Broadcast to room_b
|
||||
message_b = {"type": "room_b_message", "data": "Message for room B"}
|
||||
await websocket_service._manager.broadcast_to_room(message_b, "room_b")
|
||||
|
||||
# Verify room_b and room_both clients received, room_a did not
|
||||
for client in room_b_clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == message_b
|
||||
|
||||
for client in room_both_clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == message_b
|
||||
|
||||
for client in room_a_clients:
|
||||
assert len(client.received_messages) == 0
|
||||
|
||||
# Cleanup
|
||||
for client in room_a_clients + room_b_clients + room_both_clients:
|
||||
await client.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rapid_connect_disconnect(self, websocket_service):
|
||||
"""Test rapid connection and disconnection cycles."""
|
||||
client_id = "rapid_test_client"
|
||||
|
||||
# Perform 50 rapid connect/disconnect cycles
|
||||
for i in range(50):
|
||||
client = MockWebSocketClient(f"{client_id}_{i}", websocket_service)
|
||||
await client.connect()
|
||||
assert client.is_connected
|
||||
await client.disconnect()
|
||||
assert not client.is_connected
|
||||
|
||||
# Verify no stale connections remain
|
||||
assert len(websocket_service._manager._active_connections) == 0
|
||||
assert len(websocket_service._manager._connection_metadata) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stress_message_rate(self, websocket_service):
|
||||
"""Test high-frequency message broadcasting."""
|
||||
num_clients = 20
|
||||
num_messages = 100
|
||||
clients: List[MockWebSocketClient] = []
|
||||
|
||||
# Connect clients
|
||||
for i in range(num_clients):
|
||||
client = MockWebSocketClient(f"stress_client_{i}", websocket_service)
|
||||
await client.connect()
|
||||
clients.append(client)
|
||||
|
||||
# Send 100 messages rapidly
|
||||
for i in range(num_messages):
|
||||
message = {
|
||||
"type": "stress_test",
|
||||
"sequence": i,
|
||||
"timestamp": time.time()
|
||||
}
|
||||
await websocket_service.broadcast(message)
|
||||
|
||||
# Verify all clients received all messages
|
||||
for client in clients:
|
||||
assert len(client.received_messages) == num_messages
|
||||
# Verify messages are in order
|
||||
for i in range(num_messages):
|
||||
assert client.received_messages[i]["sequence"] == i
|
||||
|
||||
# Cleanup
|
||||
for client in clients:
|
||||
await client.disconnect()
|
||||
|
||||
|
||||
class TestWebSocketConnectionRecovery:
|
||||
"""Test WebSocket connection recovery after failures."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_recovery_after_disconnect(self, websocket_service):
|
||||
"""Test client can reconnect after unexpected disconnect."""
|
||||
client_id = "recovery_test_client"
|
||||
|
||||
# Initial connection
|
||||
client1 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client1.connect({"user_id": "test_user"})
|
||||
await client1.join_room("downloads")
|
||||
|
||||
# Simulate unexpected disconnect
|
||||
await client1.disconnect()
|
||||
assert not client1.is_connected
|
||||
|
||||
# Reconnect with same client_id
|
||||
client2 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client2.connect({"user_id": "test_user"})
|
||||
await client2.join_room("downloads")
|
||||
|
||||
# Verify new connection works
|
||||
message = {"type": "test", "data": "recovery test"}
|
||||
await websocket_service._manager.broadcast_to_room(message, "downloads")
|
||||
|
||||
assert len(client2.received_messages) == 1
|
||||
assert client2.received_messages[0] == message
|
||||
|
||||
await client2.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_room_rejoin_after_reconnection(self, websocket_service):
|
||||
"""Test client can rejoin rooms after reconnection."""
|
||||
client_id = "rejoin_test_client"
|
||||
|
||||
# Connect and join multiple rooms
|
||||
client1 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client1.connect()
|
||||
await client1.join_room("downloads")
|
||||
await client1.join_room("progress")
|
||||
await client1.join_room("updates")
|
||||
|
||||
# Verify client is in all rooms
|
||||
assert client_id in websocket_service._manager._rooms["downloads"]
|
||||
assert client_id in websocket_service._manager._rooms["progress"]
|
||||
assert client_id in websocket_service._manager._rooms["updates"]
|
||||
|
||||
# Disconnect
|
||||
await client1.disconnect()
|
||||
|
||||
# Rooms should be empty after disconnect
|
||||
for room in ["downloads", "progress", "updates"]:
|
||||
assert client_id not in websocket_service._manager._rooms.get(room, set())
|
||||
|
||||
# Reconnect and rejoin rooms
|
||||
client2 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client2.connect()
|
||||
await client2.join_room("downloads")
|
||||
await client2.join_room("progress")
|
||||
await client2.join_room("updates")
|
||||
|
||||
# Verify client is in all rooms again
|
||||
assert client_id in websocket_service._manager._rooms["downloads"]
|
||||
assert client_id in websocket_service._manager._rooms["progress"]
|
||||
assert client_id in websocket_service._manager._rooms["updates"]
|
||||
|
||||
await client2.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_delivery_after_reconnection(self, websocket_service):
|
||||
"""Test messages are delivered correctly after reconnection."""
|
||||
client_id = "delivery_test_client"
|
||||
|
||||
# Connect, receive a message, disconnect
|
||||
client1 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client1.connect()
|
||||
|
||||
message1 = {"type": "test", "sequence": 1}
|
||||
await websocket_service.broadcast(message1)
|
||||
assert len(client1.received_messages) == 1
|
||||
|
||||
await client1.disconnect()
|
||||
|
||||
# Reconnect and verify new messages are received
|
||||
client2 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client2.connect()
|
||||
|
||||
message2 = {"type": "test", "sequence": 2}
|
||||
await websocket_service.broadcast(message2)
|
||||
|
||||
# Should only receive message2 (not message1 from before disconnect)
|
||||
assert len(client2.received_messages) == 1
|
||||
assert client2.received_messages[0] == message2
|
||||
|
||||
await client2.disconnect()
|
||||
|
||||
|
||||
class TestWebSocketAuthentication:
|
||||
"""Test WebSocket authentication and token handling."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_with_authentication_metadata(
|
||||
self, websocket_service, mock_auth_token
|
||||
):
|
||||
"""Test WebSocket connection with authentication token in metadata."""
|
||||
client = MockWebSocketClient("auth_client", websocket_service)
|
||||
metadata = {
|
||||
"user_id": "test_user",
|
||||
"auth_token": mock_auth_token,
|
||||
"session_id": "session_123"
|
||||
}
|
||||
|
||||
await client.connect(metadata)
|
||||
|
||||
# Verify metadata is stored
|
||||
stored_metadata = websocket_service._manager._connection_metadata["auth_client"]
|
||||
assert stored_metadata["user_id"] == "test_user"
|
||||
assert stored_metadata["auth_token"] == mock_auth_token
|
||||
assert stored_metadata["session_id"] == "session_123"
|
||||
|
||||
await client.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_to_specific_user(self, websocket_service):
|
||||
"""Test broadcasting to specific user using metadata filtering."""
|
||||
# Connect multiple clients with different user IDs
|
||||
client1 = MockWebSocketClient("client1", websocket_service)
|
||||
await client1.connect({"user_id": "user_1"})
|
||||
|
||||
client2 = MockWebSocketClient("client2", websocket_service)
|
||||
await client2.connect({"user_id": "user_2"})
|
||||
|
||||
client3 = MockWebSocketClient("client3", websocket_service)
|
||||
await client3.connect({"user_id": "user_1"}) # Same user, different connection
|
||||
|
||||
# Broadcast to specific user
|
||||
message = {"type": "user_specific", "data": "Message for user_1"}
|
||||
|
||||
# Filter connections by user_id and send
|
||||
for conn_id, metadata in websocket_service._manager._connection_metadata.items():
|
||||
if metadata.get("user_id") == "user_1":
|
||||
ws = websocket_service._manager._active_connections[conn_id]
|
||||
await ws.send_json(message)
|
||||
|
||||
# Verify only user_1 clients received the message
|
||||
assert len(client1.received_messages) == 1
|
||||
assert client1.received_messages[0] == message
|
||||
|
||||
assert len(client3.received_messages) == 1
|
||||
assert client3.received_messages[0] == message
|
||||
|
||||
assert len(client2.received_messages) == 0
|
||||
|
||||
# Cleanup
|
||||
await client1.disconnect()
|
||||
await client2.disconnect()
|
||||
await client3.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_refresh_in_metadata(self, websocket_service):
|
||||
"""Test updating authentication token in connection metadata."""
|
||||
client = MockWebSocketClient("token_refresh_client", websocket_service)
|
||||
old_token = "old_token_12345"
|
||||
new_token = "new_token_67890"
|
||||
|
||||
# Connect with old token
|
||||
await client.connect({"user_id": "test_user", "auth_token": old_token})
|
||||
|
||||
# Verify old token is stored
|
||||
metadata = websocket_service._manager._connection_metadata["token_refresh_client"]
|
||||
assert metadata["auth_token"] == old_token
|
||||
|
||||
# Update token (simulating token refresh)
|
||||
metadata["auth_token"] = new_token
|
||||
|
||||
# Verify token is updated
|
||||
updated_metadata = websocket_service._manager._connection_metadata["token_refresh_client"]
|
||||
assert updated_metadata["auth_token"] == new_token
|
||||
|
||||
await client.disconnect()
|
||||
|
||||
|
||||
class TestWebSocketMessageOrdering:
|
||||
"""Test WebSocket message ordering guarantees."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_order_preservation(self, websocket_service):
|
||||
"""Test messages are received in the order they are sent."""
|
||||
client = MockWebSocketClient("order_test_client", websocket_service)
|
||||
await client.connect()
|
||||
|
||||
# Send 50 messages in sequence
|
||||
num_messages = 50
|
||||
for i in range(num_messages):
|
||||
message = {
|
||||
"type": "sequence_test",
|
||||
"sequence": i,
|
||||
"timestamp": time.time()
|
||||
}
|
||||
await websocket_service.broadcast(message)
|
||||
|
||||
# Verify all messages received in order
|
||||
assert len(client.received_messages) == num_messages
|
||||
for i in range(num_messages):
|
||||
assert client.received_messages[i]["sequence"] == i
|
||||
|
||||
await client.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_broadcast_order(self, websocket_service):
|
||||
"""Test message ordering with concurrent broadcasts to different rooms."""
|
||||
# Create clients in two rooms
|
||||
room1_client = MockWebSocketClient("room1_client", websocket_service)
|
||||
await room1_client.connect()
|
||||
await room1_client.join_room("room1")
|
||||
|
||||
room2_client = MockWebSocketClient("room2_client", websocket_service)
|
||||
await room2_client.connect()
|
||||
await room2_client.join_room("room2")
|
||||
|
||||
both_rooms_client = MockWebSocketClient("both_client", websocket_service)
|
||||
await both_rooms_client.connect()
|
||||
await both_rooms_client.join_room("room1")
|
||||
await both_rooms_client.join_room("room2")
|
||||
|
||||
# Send interleaved messages to both rooms
|
||||
for i in range(10):
|
||||
message1 = {"type": "room1_msg", "sequence": i}
|
||||
await websocket_service._manager.broadcast_to_room(message1, "room1")
|
||||
|
||||
message2 = {"type": "room2_msg", "sequence": i}
|
||||
await websocket_service._manager.broadcast_to_room(message2, "room2")
|
||||
|
||||
# Verify room1_client received only room1 messages in order
|
||||
assert len(room1_client.received_messages) == 10
|
||||
for i in range(10):
|
||||
assert room1_client.received_messages[i]["type"] == "room1_msg"
|
||||
assert room1_client.received_messages[i]["sequence"] == i
|
||||
|
||||
# Verify room2_client received only room2 messages in order
|
||||
assert len(room2_client.received_messages) == 10
|
||||
for i in range(10):
|
||||
assert room2_client.received_messages[i]["type"] == "room2_msg"
|
||||
assert room2_client.received_messages[i]["sequence"] == i
|
||||
|
||||
# Verify both_rooms_client received all messages (may be interleaved)
|
||||
assert len(both_rooms_client.received_messages) == 20
|
||||
room1_msgs = [msg for msg in both_rooms_client.received_messages if msg["type"] == "room1_msg"]
|
||||
room2_msgs = [msg for msg in both_rooms_client.received_messages if msg["type"] == "room2_msg"]
|
||||
assert len(room1_msgs) == 10
|
||||
assert len(room2_msgs) == 10
|
||||
|
||||
# Cleanup
|
||||
await room1_client.disconnect()
|
||||
await room2_client.disconnect()
|
||||
await both_rooms_client.disconnect()
|
||||
|
||||
|
||||
class TestWebSocketBroadcastFiltering:
|
||||
"""Test WebSocket broadcast filtering to specific clients."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_to_all_except_sender(self, websocket_service):
|
||||
"""Test broadcasting to all clients except the sender."""
|
||||
# Connect multiple clients
|
||||
sender = MockWebSocketClient("sender", websocket_service)
|
||||
await sender.connect()
|
||||
|
||||
clients = []
|
||||
for i in range(5):
|
||||
client = MockWebSocketClient(f"client_{i}", websocket_service)
|
||||
await client.connect()
|
||||
clients.append(client)
|
||||
|
||||
# Broadcast to all except sender
|
||||
message = {"type": "broadcast", "data": "Message to all except sender"}
|
||||
|
||||
for conn_id in websocket_service._manager._active_connections:
|
||||
if conn_id != "sender":
|
||||
ws = websocket_service._manager._active_connections[conn_id]
|
||||
await ws.send_json(message)
|
||||
|
||||
# Verify sender did not receive message
|
||||
assert len(sender.received_messages) == 0
|
||||
|
||||
# Verify all other clients received message
|
||||
for client in clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == message
|
||||
|
||||
# Cleanup
|
||||
await sender.disconnect()
|
||||
for client in clients:
|
||||
await client.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_filtered_by_metadata(self, websocket_service):
|
||||
"""Test broadcasting filtered by connection metadata."""
|
||||
# Connect clients with different roles
|
||||
admin_clients = []
|
||||
for i in range(3):
|
||||
client = MockWebSocketClient(f"admin_{i}", websocket_service)
|
||||
await client.connect({"role": "admin", "user_id": f"admin_{i}"})
|
||||
admin_clients.append(client)
|
||||
|
||||
user_clients = []
|
||||
for i in range(3):
|
||||
client = MockWebSocketClient(f"user_{i}", websocket_service)
|
||||
await client.connect({"role": "user", "user_id": f"user_{i}"})
|
||||
user_clients.append(client)
|
||||
|
||||
# Broadcast only to admins
|
||||
admin_message = {"type": "admin_only", "data": "Admin notification"}
|
||||
|
||||
for conn_id, metadata in websocket_service._manager._connection_metadata.items():
|
||||
if metadata.get("role") == "admin":
|
||||
ws = websocket_service._manager._active_connections[conn_id]
|
||||
await ws.send_json(admin_message)
|
||||
|
||||
# Verify only admin clients received message
|
||||
for client in admin_clients:
|
||||
assert len(client.received_messages) == 1
|
||||
assert client.received_messages[0] == admin_message
|
||||
|
||||
for client in user_clients:
|
||||
assert len(client.received_messages) == 0
|
||||
|
||||
# Cleanup
|
||||
for client in admin_clients + user_clients:
|
||||
await client.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_room_based_filtering(self, websocket_service):
|
||||
"""Test combining room membership and metadata filtering."""
|
||||
# Create clients with different metadata in the same room
|
||||
premium_client = MockWebSocketClient("premium", websocket_service)
|
||||
await premium_client.connect({"subscription": "premium"})
|
||||
await premium_client.join_room("downloads")
|
||||
|
||||
free_client = MockWebSocketClient("free", websocket_service)
|
||||
await free_client.connect({"subscription": "free"})
|
||||
await free_client.join_room("downloads")
|
||||
|
||||
# Send premium-only message to downloads room
|
||||
premium_message = {"type": "premium_feature", "data": "Premium notification"}
|
||||
|
||||
# Get clients in downloads room with premium subscription
|
||||
room_members = websocket_service._manager._rooms.get("downloads", set())
|
||||
for conn_id in room_members:
|
||||
metadata = websocket_service._manager._connection_metadata.get(conn_id, {})
|
||||
if metadata.get("subscription") == "premium":
|
||||
ws = websocket_service._manager._active_connections[conn_id]
|
||||
await ws.send_json(premium_message)
|
||||
|
||||
# Verify only premium client received message
|
||||
assert len(premium_client.received_messages) == 1
|
||||
assert premium_client.received_messages[0] == premium_message
|
||||
|
||||
assert len(free_client.received_messages) == 0
|
||||
|
||||
# Cleanup
|
||||
await premium_client.disconnect()
|
||||
await free_client.disconnect()
|
||||
|
||||
|
||||
class TestWebSocketEdgeCases:
|
||||
"""Test WebSocket edge cases and error conditions."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_connection_ids(self, websocket_service):
|
||||
"""Test handling duplicate connection IDs (should replace old connection)."""
|
||||
client_id = "duplicate_id"
|
||||
|
||||
# First connection
|
||||
client1 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client1.connect()
|
||||
|
||||
# Send message to first connection
|
||||
message1 = {"type": "test", "sequence": 1}
|
||||
await websocket_service.broadcast(message1)
|
||||
assert len(client1.received_messages) == 1
|
||||
|
||||
# Second connection with same ID (should replace first)
|
||||
client2 = MockWebSocketClient(client_id, websocket_service)
|
||||
await client2.connect()
|
||||
|
||||
# Only one connection should exist
|
||||
assert len(websocket_service._manager._active_connections) == 1
|
||||
|
||||
# Send message to second connection
|
||||
message2 = {"type": "test", "sequence": 2}
|
||||
await websocket_service.broadcast(message2)
|
||||
|
||||
# Second client should receive message
|
||||
assert len(client2.received_messages) == 1
|
||||
assert client2.received_messages[0] == message2
|
||||
|
||||
await client2.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_leave_nonexistent_room(self, websocket_service):
|
||||
"""Test leaving a room that doesn't exist or client isn't in."""
|
||||
client = MockWebSocketClient("test_client", websocket_service)
|
||||
await client.connect()
|
||||
|
||||
# Should not raise error
|
||||
await client.leave_room("nonexistent_room")
|
||||
|
||||
await client.disconnect()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_to_disconnected_client(self, websocket_service):
|
||||
"""Test sending message to a client that has disconnected."""
|
||||
client = MockWebSocketClient("disconnect_test", websocket_service)
|
||||
await client.connect()
|
||||
|
||||
# Disconnect
|
||||
await client.disconnect()
|
||||
|
||||
# Attempt to broadcast (should not raise error)
|
||||
message = {"type": "test", "data": "test"}
|
||||
await websocket_service.broadcast(message)
|
||||
|
||||
# Client should not receive message (already disconnected)
|
||||
assert len(client.received_messages) == 0
|
||||
Reference in New Issue
Block a user