# conftest.py import asyncio import pytest import pytest_asyncio from unittest.mock import AsyncMock, MagicMock, patch import psycopg2 from fastapi import FastAPI from fastapi.testclient import TestClient # Import the modules we need to test # Note: These would be the actual file names from the implementation import sys from pathlib import Path # Mock the actual implementations for testing @pytest.fixture def mock_postgresql_pool(): """Mock PostgreSQL connection pool""" pool = MagicMock() conn = MagicMock() cursor = MagicMock() cursor.fetchone.return_value = None cursor.fetchall.return_value = [] conn.cursor.return_value = cursor conn.__enter__ = MagicMock(return_value=conn) conn.__exit__ = MagicMock(return_value=False) pool.getconn.return_value = conn pool.putconn.return_value = None return pool, conn, cursor @pytest.fixture def sample_directive(): """Sample directive for testing""" return { "id": "cmd-12345", "type": "directive", "action": "deploy", "target": "production", "parameters": {"version": "1.2.3"}, "timestamp": "2024-01-15T10:30:00Z", "priority": "high" } @pytest.fixture def sample_heartbeat(): """Sample heartbeat message""" return { "type": "heartbeat", "timestamp": "2024-01-15T10:30:30Z" } @pytest.fixture def valid_api_key(): return "test-api-key-12345" @pytest_asyncio.fixture async def websocket_client_mock(): """Mock WebSocket client""" ws = AsyncMock() ws.send_json = AsyncMock() ws.receive_json = AsyncMock() ws.receive_text = AsyncMock() ws.close = AsyncMock() ws.accept = AsyncMock() return ws # tests/test_connection_manager.py import pytest import asyncio from unittest.mock import AsyncMock, MagicMock, patch from typing import Dict, Set # Mock ConnectionManager for testing class ConnectionManager: def __init__(self): self.active_connections: Dict[str, AsyncMock] = {} self.client_metadata: Dict[str, dict] = {} async def connect(self, client_id: str, websocket: AsyncMock, metadata: dict = None): await websocket.accept() self.active_connections[client_id] = websocket self.client_metadata[client_id] = metadata or {} def disconnect(self, client_id: str): if client_id in self.active_connections: del self.active_connections[client_id] del self.client_metadata[client_id] async def broadcast(self, message: dict): disconnected = [] for client_id, connection in self.active_connections.items(): try: await connection.send_json(message) except Exception: disconnected.append(client_id) for client_id in disconnected: self.disconnect(client_id) async def send_to_client(self, client_id: str, message: dict): if client_id in self.active_connections: await self.active_connections[client_id].send_json(message) def get_client_count(self): return len(self.active_connections) def get_clients(self): return list(self.active_connections.keys()) @pytest.mark.asyncio class TestConnectionManager: async def test_connect_adds_client(self): manager = ConnectionManager() ws = AsyncMock() await manager.connect("client-1", ws, {"role": "claude"}) assert manager.get_client_count() == 1 assert "client-1" in manager.get_clients() ws.accept.assert_called_once() async def test_connect_with_metadata(self): manager = ConnectionManager() ws = AsyncMock() metadata = {"role": "claude", "version": "1.0"} await manager.connect("client-1", ws, metadata) assert manager.client_metadata["client-1"] == metadata async def test_disconnect_removes_client(self): manager = ConnectionManager() ws = AsyncMock() await manager.connect("client-1", ws) manager.disconnect("client-1") assert manager.get_client_count() == 0 assert "client-1" not in manager.get_clients() async def test_broadcast_sends_to_all(self): manager = ConnectionManager() ws1 = AsyncMock() ws2 = AsyncMock() await manager.connect("client-1", ws1) await manager.connect("client-2", ws2) message = {"type": "directive", "data": "test"} await manager.broadcast(message) ws1.send_json.assert_called_once_with(message) ws2.send_json.assert_called_once_with(message) async def test_broadcast_handles_failure(self): manager = ConnectionManager() ws1 = AsyncMock() ws2 = AsyncMock() ws1.send_json = AsyncMock(side_effect=Exception("Connection lost")) await manager.connect("client-1", ws1) await manager.connect("client-2", ws2) message = {"type": "directive", "data": "test"} await manager.broadcast(message) # client-1 should be disconnected due to failure assert "client-1" not in manager.get_clients() assert "client-2" in manager.get_clients() async def test_send_to_client_specific(self): manager = ConnectionManager() ws1 = AsyncMock() ws2 = AsyncMock() await manager.connect("client-1", ws1) await manager.connect("client-2", ws2) message = {"type": "directive", "data": "private"} await manager.send_to_client("client-1", message) ws1.send_json.assert_called_once_with(message) ws2.send_json.assert_not_called() async def test_send_to_nonexistent_client(self): manager = ConnectionManager() # Should not raise exception await manager.send_to_client("ghost-client", {"type": "test"}) # tests/test_bridge_ws.py import pytest import pytest_asyncio import json from unittest.mock import AsyncMock, MagicMock, patch, Mock from fastapi import FastAPI, WebSocket, WebSocketDisconnect, status from fastapi.testclient import TestClient import asyncio # Mock the WebSocket endpoint implementation class BridgeWebSocketServer: def __init__(self, connection_manager, api_key): self.manager = connection_manager self.api_key = api_key async def handle_connection(self, websocket: WebSocket, token: str = None): # Authenticate if token != self.api_key: await websocket.close(code=status.WS_1008_POLICY_VIOLATION) return client_id = f"client-{id(websocket)}" await self.manager.connect(client_id, websocket, {"connected_at": "2024-01-15"}) try: while True: data = await websocket.receive_json() await self._handle_message(client_id, data) except WebSocketDisconnect: self.manager.disconnect(client_id) async def _handle_message(self, client_id: str, data: dict): msg_type = data.get("type") if msg_type == "ack": await self._handle_ack(client_id, data) elif msg_type == "response": await self._handle_response(client_id, data) elif msg_type == "heartbeat": await self._handle_heartbeat(client_id, data) async def _handle_ack(self, client_id: str, data: dict): # Update database or log acknowledgment pass async def _handle_response(self, client_id: str, data: dict): # Store response in database pass async def _handle_heartbeat(self, client_id: str, data: dict): # Send heartbeat back pass @pytest.mark.asyncio class TestBridgeWebSocketServer: async def test_authentication_with_valid_token(self): manager = MagicMock() manager.connect = AsyncMock() server = BridgeWebSocketServer(manager, "valid-key") ws = AsyncMock() ws.receive_json = AsyncMock(side_effect=WebSocketDisconnect()) await server.handle_connection(ws, token="valid-key") ws.close.assert_not_called() manager.connect.assert_called_once() async def test_authentication_with_invalid_token(self): manager = MagicMock() server = BridgeWebSocketServer(manager, "valid-key") ws = AsyncMock() await server.handle_connection(ws, token="invalid-key") ws.close.assert_called_once() assert ws.close.call_args[1]["code"] == status.WS_1008_POLICY_VIOLATION async def test_message_handling_ack(self): manager = MagicMock() server = BridgeWebSocketServer(manager, "key") with patch.object(server, '_handle_ack', new=AsyncMock()) as mock_ack: await server._handle_message("client-1", {"type": "ack", "id": "123"}) mock_ack.assert_called_once() async def test_message_handling_response(self): manager = MagicMock() server = BridgeWebSocketServer(manager, "key") with patch.object(server, '_handle_response', new=AsyncMock()) as mock_resp: await server._handle_message("client-1", {"type": "response", "data": "test"}) mock_resp.assert_called_once() async def test_disconnect_handling(self): manager = MagicMock() manager.disconnect = MagicMock() server = BridgeWebSocketServer(manager, "key") ws = AsyncMock() ws.receive_json = AsyncMock(side_effect=WebSocketDisconnect()) await server.handle_connection(ws, token="key") manager.disconnect.assert_called_once() # tests/test_bridge_ws_client.py import pytest import pytest_asyncio import json import asyncio from unittest.mock import AsyncMock, MagicMock, patch, call from websockets.exceptions import ConnectionClosed, InvalidStatusCode class BridgeWebSocketClient: def __init__(self, uri: str, api_key: str, fallback_poll_interval: int = 5): self.uri = uri self.api_key = api_key self.fallback_poll_interval = fallback_poll_interval self.ws = None self.connected = False self.reconnect_attempts = 0 self.max_reconnect_attempts = 5 self.backoff_base = 1.5 self.message_handlers = { "directive": self._on_directive, "heartbeat": self._on_heartbeat, "status": self._on_status } self.directives_received = [] async def connect(self): while self.reconnect_attempts < self.max_reconnect_attempts: try: import websockets self.ws = await websockets.connect( f"{self.uri}?token={self.api_key}" ) self.connected = True self.reconnect_attempts = 0 asyncio.create_task(self._receive_loop()) asyncio.create_task(self._heartbeat_loop()) return True except Exception as e: self.reconnect_attempts += 1 wait_time = self.backoff_base ** self.reconnect_attempts await asyncio.sleep(wait_time) return False async def disconnect(self): self.connected = False if self.ws: await self.ws.close() async def _receive_loop(self): try: while self.connected: message = await self.ws.recv() data = json.loads(message) await self._handle_message(data) except ConnectionClosed: self.connected = False await self._attempt_reconnect() async def _handle_message(self, data: dict): msg_type = data.get("type") handler = self.message_handlers.get(msg_type) if handler: await handler(data) async def _on_directive(self, data: dict): self.directives_received.append(data) # Send ack await self.send_message({"type": "ack", "id": data.get("id")}) async def _on_heartbeat(self, data: dict): # Respond to heartbeat await self.send_message({"type": "heartbeat", "timestamp": "2024-01-15T10:30:00Z"}) async def _on_status(self, data: dict): pass async def send_message(self, message: dict): if self.connected and self.ws: await self.ws.send(json.dumps(message)) async def _heartbeat_loop(self): while self.connected: await asyncio.sleep(30) if self.connected: await self.send_message({"type": "heartbeat", "timestamp": "2024-01-15T10:30:00Z"}) async def _attempt_reconnect(self): self.connected = False await self.connect() @pytest.mark.asyncio class TestBridgeWebSocketClient: @pytest_asyncio.fixture async def client(self): ws_client = BridgeWebSocketClient( uri="ws://localhost:8000/bridge/ws", api_key="test-key" ) yield ws_client await ws_client.disconnect() async def test_initialization(self): client = BridgeWebSocketClient("ws://test", "key") assert client.uri == "ws://test" assert client.api_key == "key" assert client.connected is False assert client.reconnect_attempts == 0 async def test_connect_success(self, client): with patch('websockets.connect', new_callable=AsyncMock) as mock_connect: mock_ws = AsyncMock() mock_connect.return_value = mock_ws result = await client.connect() assert result is True assert client.connected is True assert client.ws == mock_ws mock_connect.assert_called_once() async def test_connect_with_auth_token(self, client): with patch('websockets.connect', new_callable=AsyncMock) as mock_connect: mock_ws = AsyncMock() mock_connect.return_value = mock_ws await client.connect() # Check that token was included in URI call_args = mock_connect.call_args[0][0] assert "token=test-key" in call_args async def test_connect_failure_retry(self, client): with patch('websockets.connect', new_callable=AsyncMock) as mock_connect: mock_connect.side_effect = Exception("Connection refused") with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep: result = await client.connect() assert result is False assert client.reconnect_attempts == 5 # max attempts assert mock_connect.call_count == 5 async def test_receive_directive(self, client): directive = { "type": "directive", "id": "cmd-123", "action": "test" } with patch.object(client, '_on_directive', new_callable=AsyncMock) as mock_handler: await client._handle_message(directive) mock_handler.assert_called_once_with(directive) async def test_on_directive_sends_ack(self, client): directive = { "type": "directive", "id": "cmd-123", "action": "test" } client.ws = AsyncMock() client.connected = True await client._on_directive(directive) # Check that ack was sent client.ws.send.assert_called_once() sent_data = json.loads(client.ws.send.call_args[0][0]) assert sent_data["type"] == "ack" assert sent_data["id"] == "cmd-123" async def test_heartbeat_response(self, client): client.ws = AsyncMock() client.connected = True heartbeat = {"type": "heartbeat", "timestamp": "2024-01-15T10:00:00Z"} await client._on_heartbeat(heartbeat) # Should send heartbeat back assert client.ws.send.called async def test_disconnect_cleanup(self, client): client.ws = AsyncMock() client.connected = True await client.disconnect() assert client.connected is False client.ws.close.assert_called_once() async def test_reconnect_logic(self, client): with patch.object(client, 'connect', new_callable=AsyncMock) as mock_connect: await client._attempt_reconnect() mock_connect.assert_called_once() async def test_receive_loop_handles_connection_closed(self, client): client.ws = AsyncMock() client.ws.recv = AsyncMock(side_effect=Exception("Connection closed")) client.connected = True with patch.object(client, '_attempt_reconnect', new_callable=AsyncMock) as mock_reconnect: await client._receive_loop() mock_reconnect.assert_called_once() async def test_malformed_json_handling(self, client): client.ws = AsyncMock() client.ws.recv = AsyncMock(return_value="invalid json{") client.connected = True # Should not raise exception with pytest.raises(Exception): await client._receive_loop() # tests/test_postgres_notify.py import pytest import pytest_asyncio from unittest.mock import MagicMock, patch, AsyncMock import psycopg2 import select class PostgresNotifyListener: def __init__(self, dsn: str): self.dsn = dsn self.conn = None self.listening = False self.callbacks = [] def connect(self): self.conn = psycopg2.connect(self.dsn) self.conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT) def listen(self, channel: str): if not self.conn: self.connect() cursor = self.conn.cursor() cursor.execute(f"LISTEN {channel};") self.listening = True def add_callback(self, callback): self.callbacks.append(callback) async def poll(self, timeout: float = 1.0): if not self.conn: return if select.select([self.conn], [], [], timeout) == ([], [], []): return self.conn.poll() while self.conn.notifies: notify = self.conn.notifies.pop(0) for callback in self.callbacks: if asyncio.iscoroutinefunction(callback): await callback(notify) else: callback(notify) def close(self): if self.conn: self.conn.close() self.listening = False @pytest.mark.asyncio class TestPostgresNotify: @pytest.fixture def listener(self): return PostgresNotifyListener("postgresql://user:pass@localhost/db") def test_initialization(self, listener): assert listener.dsn == "postgresql://user:pass@localhost/db" assert listener.conn is None assert listener.listening is False assert listener.callbacks == [] @patch('psycopg2.connect') def test_connect(self, mock_connect, listener): mock_conn = MagicMock() mock_connect.return_value = mock_conn listener.connect() mock_connect.assert_called_once_with(listener.dsn) mock_conn.set_isolation_level.assert_called_once() assert listener.conn == mock_conn @patch('psycopg2.connect') def test_listen(self, mock_connect, listener): mock_conn = MagicMock() mock_cursor = MagicMock() mock_conn.cursor.return_value = mock_cursor mock_connect.return_value = mock_conn listener.listen("command_queue") mock_cursor.execute.assert_called_once_with("LISTEN command_queue;") assert listener.listening is True def test_add_callback(self, listener): callback = MagicMock() listener.add_callback(callback) assert callback in listener.callbacks @patch('select.select') @patch('psycopg2.connect') async def test_poll_with_notification(self, mock_connect, mock_select, listener): mock_conn = MagicMock() mock_notify = MagicMock() mock_notify.channel = "command_queue" mock_notify.payload = '{"id": "123"}' mock_conn.notifies = [mock_notify] mock_conn.poll = MagicMock() mock_connect.return_value = mock_conn listener.conn = mock_conn mock_select.return_value = ([mock_conn], [], []) callback = AsyncMock() listener.add_callback(callback) await listener.poll() callback.assert_called_once_with(mock_notify) @patch('select.select') async def test_poll_timeout(self, mock_select, listener): mock_select.return_value = ([], [], []) listener.conn = MagicMock() # Should not block indefinitely await listener.poll(timeout=0.1) def test_close(self, listener): mock_conn = MagicMock() listener.conn = mock_conn listener.listening = True listener.close() mock_conn.close.assert_called_once() assert listener.listening is False # tests/test_bridge_ws_integration.py import pytest import pytest_asyncio import json from unittest.mock import AsyncMock, MagicMock, patch, Mock from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, Header, HTTPException, status from fastapi.testclient import TestClient from starlette.testclient import TestClient as StarletteTestClient from starlette.websockets import WebSocketTestSession import asyncio # Create test FastAPI app def create_test_app(): app = FastAPI() manager = MagicMock() manager.active_connections = {} manager.connect = AsyncMock() manager.disconnect = MagicMock() manager.broadcast = AsyncMock() manager.send_to_client = AsyncMock() @app.websocket("/bridge/ws") async def websocket_endpoint(websocket: WebSocket, token: str = None): if token != "valid-token": await websocket.close(code=status.WS_1008_POLICY_VIOLATION) return client_id = f"client-{id(websocket)}" await manager.connect(client_id, websocket, {"token": token}) try: while True: data = await websocket.receive_json() if data.get("type") == "ack": await manager.send_to_client(client_id, {"type": "ack_received"}) except WebSocketDisconnect: manager.disconnect(client_id) return app, manager @pytest.fixture def test_app(): return create_test_app() @pytest.mark.asyncio class TestWebSocketIntegration: def test_websocket_auth_valid_token(self, test_app): app, manager = test_app client = TestClient(app) with client.websocket_connect("/bridge/ws?token=valid-token") as websocket: # Should stay connected pass manager.connect.assert_called_once() def test_websocket_auth_invalid_token(self, test_app): app, manager = test_app client = TestClient(app) with client.websocket_connect("/bridge/ws?token=invalid-token") as websocket: # Should receive close pass def test_websocket_message_echo(self, test_app): app, manager = test_app client = TestClient(app) with client.websocket_connect("/bridge/ws?token=valid-token") as websocket: websocket.send_json({"type": "ack", "id": "123"}) # In real implementation, would receive response def test_websocket_disconnect_cleanup(self, test_app): app, manager = test_app client = TestClient(app) with client.websocket_connect("/bridge/ws?token=valid-token") as websocket: pass manager.disconnect.assert_called_once() # tests/test_bridge_ws_client_unit.py import pytest import pytest_asyncio import json import asyncio from unittest.mock import AsyncMock, MagicMock, patch, call from datetime import datetime class TestBridgeWebSocketClientUnit: @pytest_asyncio.fixture async def ws_client(self): client = BridgeWebSocketClient( uri="ws://localhost:8000/bridge/ws", api_key="test-key" ) yield client if client.connected: await client.disconnect() async def test_client_initialization(self): client = BridgeWebSocketClient("ws://test", "key", fallback_poll_interval=10) assert client.uri == "ws://test" assert client.api_key == "key" assert client.fallback_poll_interval == 10 assert client.max_reconnect_attempts == 5 async def test_connect_success(self, ws_client): with patch('websockets.connect', new_callable=AsyncMock) as mock_ws: mock_socket = AsyncMock() mock_ws.return_value = mock_socket result = await ws_client.connect() assert result is True assert ws_client.connected is True assert ws_client.ws == mock_socket mock_ws.assert_called_with("ws://localhost:8000/bridge/ws?token=test-key") async def test_connect_failure_with_retry(self, ws_client): with patch('websockets.connect', new_callable=AsyncMock) as mock_ws: mock_ws.side_effect = Exception("Connection refused") with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep: result = await ws_client.connect() assert result is False assert ws_client.connected is False assert mock_ws.call_count == 5 # max_reconnect_attempts async def test_exponential_backoff(self, ws_client): with patch('websockets.connect', new_callable=AsyncMock) as mock_ws: mock_ws.side_effect = Exception("Connection refused") sleep_times = [] async def capture_sleep(t): sleep_times.append(t) with patch('asyncio.sleep', side_effect=capture_sleep): try: await ws_client.connect() except: pass # Check exponential backoff: 1.5^1, 1.5^2, etc assert len(sleep_times) > 0 assert sleep_times[0] == 1.5 ** 1 assert sleep_times[1] == 1.5 ** 2 async def test_receive_directive(self, ws_client): directive = { "type": "directive", "id": "cmd-123", "action": "deploy" } ws_client.ws = AsyncMock() ws_client.ws.recv = AsyncMock(return_value=json.dumps(directive)) ws_client.connected = True # Mock the reconnect to stop the loop after one iteration ws_client.ws.recv.side_effect = [json.dumps(directive), Exception("Stop")] with patch.object(ws_client, '_attempt_reconnect', new_callable=AsyncMock): try: await ws_client._receive_loop() except: pass assert len(ws_client.directives_received) == 1 assert ws_client.directives_received[0]["id"] == "cmd-123" async def test_send_ack_on_directive(self, ws_client): directive = { "type": "directive", "id": "cmd-123", "action": "deploy" } ws_client.ws = AsyncMock() ws_client.connected = True await ws_client._on_directive(directive) ws_client.ws.send.assert_called_once() sent_msg = json.loads(ws_client.ws.send.call_args[0][0]) assert sent_msg["type"] == "ack" assert sent_msg["id"] == "cmd-123" async def test_heartbeat_response(self, ws_client): ws_client.ws = AsyncMock() ws_client.connected = True heartbeat = {"type": "heartbeat", "timestamp": "2024-01-15T10:00:00Z"} await ws_client._on_heartbeat(heartbeat) assert ws_client.ws.send.called async def test_malformed_json_handling(self, ws_client): ws_client.ws = AsyncMock() ws_client.ws.recv = AsyncMock(return_value="not valid json") ws_client.connected = True # Should raise exception that triggers reconnect with pytest.raises(Exception): await ws_client._receive_loop() async def test_disconnect_cleanup(self, ws_client): ws_client.ws = AsyncMock() ws_client.connected = True await ws_client.disconnect() assert ws_client.connected is False ws_client.ws.close.assert_called_once() async def test_fallback_to_polling(self, ws_client): # Simulate WebSocket failure with patch.object(ws_client, 'connect', return_value=False): with patch.object(ws_client, '_poll_fallback', new_callable=AsyncMock) as mock_poll: # Trigger fallback await ws_client.connect_with_fallback() mock_poll.assert_called_once() async def test_connection_manager_integration(self, ws_client): # Test that client properly registers with connection manager manager = MagicMock() with patch('websockets.connect', new_callable=AsyncMock) as mock_ws: mock_socket = AsyncMock() mock_ws.return_value = mock_socket await ws_client.connect() # Verify connection state assert ws_client.ws is not None # tests/test_postgres_listen_notify.py import pytest import pytest_asyncio from unittest.mock import MagicMock, patch, AsyncMock, call import psycopg2 import select import json import asyncio class TestPostgresListenNotify: @pytest.fixture def db_config(self): return { "host": "localhost", "port": 25432, "database": "postgres", "user": "postgres", "password": "test-pass" } @patch('psycopg2.connect') def test_listen_setup(self, mock_connect, db_config): mock_conn = MagicMock() mock_cursor = MagicMock() mock_conn.cursor.return_value = mock_cursor mock_connect.return_value = mock_conn listener = PostgresNotifyListener( f"postgresql://{db_config['user']}:{db_config['password']}@" f"{db_config['host']}:{db_config['port']}/{db_config['database']}" ) listener.connect() listener.listen("command_queue") mock_cursor.execute.assert_called_with("LISTEN command_queue;") @patch('select.select') @patch('psycopg2.connect') async def test_notification_received(self, mock_connect, mock_select, db_config): mock_conn = MagicMock() mock_notify = MagicMock() mock_notify.channel = "command_queue" mock_notify.payload = '{"id": "cmd-123", "action": "deploy"}' mock_conn.notifies = [mock_notify] mock_conn.poll = MagicMock() mock_connect.return_value = mock_conn mock_select.return_value = ([mock_conn], [], []) listener = PostgresNotifyListener("postgresql://user:pass@localhost/db") listener.conn = mock_conn callback_called = False received_payload = None def callback(notify): nonlocal callback_called, received_payload callback_called = True received_payload = notify.payload listener.add_callback(callback) await listener.poll(timeout=1.0) assert callback_called is True assert received_payload == '{"id": "cmd-123", "action": "deploy"}' @patch('select.select') async def test_poll_timeout(self, mock_select): mock_select.return_value = ([], [], []) listener = PostgresNotifyListener("postgresql://user:pass@localhost/db") listener.conn = MagicMock() # Should return without calling callbacks callback = MagicMock() listener.add_callback(callback) await listener.poll(timeout=0.1) callback.assert_not_called() def test_trigger_creation(self): # Test that the trigger SQL is correct trigger_sql = """ CREATE OR REPLACE FUNCTION notify_command_queue() RETURNS TRIGGER AS $$ BEGIN PERFORM pg_notify('command_queue', row_to_json(NEW)::text); RETURN NEW; END; $$ LANGUAGE plpgsql; CREATE TRIGGER command_queue_trigger AFTER INSERT ON genesis_bridge.command_queue FOR EACH ROW EXECUTE FUNCTION notify_command_queue(); """ assert "pg_notify" in trigger_sql assert "command_queue" in trigger_sql assert "AFTER INSERT" in trigger_sql @patch('psycopg2.connect') async def test_multiple_callbacks(self, mock_connect): listener = PostgresNotifyListener("postgresql://user:pass@localhost/db") cb1 = MagicMock() cb2 = MagicMock() listener.add_callback(cb1) listener.add_callback(cb2) # Simulate notification mock_notify = MagicMock() mock_notify.payload = "{}" for cb in listener.callbacks: cb(mock_notify) cb1.assert_called_once() cb2.assert_called_once() # tests/test_websocket_integration.py import pytest import pytest_asyncio import json import asyncio from unittest.mock import AsyncMock, MagicMock, patch, Mock from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.testclient import TestClient import websockets # Integration tests for full WebSocket flow @pytest.mark.asyncio class TestWebSocketIntegration: @pytest_asyncio.fixture async def full_app(self): """Create full FastAPI app with WebSocket endpoint""" app = FastAPI() manager = ConnectionManager() @app.websocket("/bridge/ws") async def websocket_endpoint(websocket: WebSocket, token: str = None): if token != "valid-token": await websocket.close(code=1008) return client_id = f"client-{id(websocket)}" await manager.connect(client_id, websocket, {"token": token}) try: while True: data = await websocket.receive_json() if data.get("type") == "ack": await manager.send_to_client(client_id, {"type": "ack_received"}) except WebSocketDisconnect: manager.disconnect(client_id) return app, manager def test_full_connection_flow(self, full_app): app, manager = full_app client = TestClient(app) with client.websocket_connect("/bridge/ws?token=valid-token") as websocket: # Send a message websocket.send_json({"type": "ack", "id": "123"}) # Should receive response # Note: In real test, we'd need to handle the async nature properly def test_auth_failure(self, full_app): app, manager = full_app client = TestClient(app) with client.websocket_connect("/bridge/ws?token=invalid-token") as websocket: # Should be closed immediately pass @pytest.mark.asyncio async def test_broadcast_to_multiple_clients(self): manager = ConnectionManager() ws1 = AsyncMock() ws2 = AsyncMock() ws3 = AsyncMock() await manager.connect("client-1", ws1) await manager.connect("client-2", ws2) await manager.connect("client-3", ws3) message = {"type": "directive", "data": "broadcast"} await manager.broadcast(message) ws1.send_json.assert_called_with(message) ws2.send_json.assert_called_with(message) ws3.send_json.assert_called_with(message) @pytest.mark.asyncio async def test_targeted_message(self): manager = ConnectionManager() ws1 = AsyncMock() ws2 = AsyncMock() await manager.connect("client-1", ws1) await manager.connect("client-2", ws2) message = {"type": "directive", "data": "private"} await manager.send_to_client("client-1", message) ws1.send_json.assert_called_once_with(message) ws2.send_json.assert_not_called() @pytest.mark.asyncio async def test_heartbeat_mechanism(self): client = BridgeWebSocketClient("ws://test", "key") client.ws = AsyncMock() client.connected = True # Mock asyncio.sleep to speed up test sleep_calls = [] async def mock_sleep(t): sleep_calls.append(t) if len(sleep_calls) >= 3: client.connected = False # Stop after 3 heartbeats with patch('asyncio.sleep', side_effect=mock_sleep): await client._heartbeat_loop() # Should have sent 3 heartbeats assert client.ws.send.call_count == 3 # Each heartbeat should be 30 seconds apart assert all(s == 30 for s in sleep_calls) # tests/test_edge_cases.py import pytest import pytest_asyncio import json from unittest.mock import AsyncMock, MagicMock, patch, Mock import asyncio @pytest.mark.asyncio class TestEdgeCases: async def test_rapid_connect_disconnect(self): """Test rapid connect/disconnect cycles don't leak resources""" manager = ConnectionManager() for i in range(100): ws = AsyncMock() await manager.connect(f"client-{i}", ws) manager.disconnect(f"client-{i}") assert manager.get_client_count() == 0 async def test_broadcast_with_failing_clients(self): """Test broadcast continues even if some clients fail""" manager = ConnectionManager() ws_good = AsyncMock() ws_bad = AsyncMock() ws_bad.send_json = AsyncMock(side_effect=Exception("Connection lost")) await manager.connect("good", ws_good) await manager.connect("bad", ws_bad) # Should not raise exception await manager.broadcast({"type": "test"}) # Good client should still receive message ws_good.send_json.assert_called_once() async def test_large_message_handling(self): """Test handling of large JSON payloads""" client = BridgeWebSocketClient("ws://test", "key") client.ws = AsyncMock() client.connected = True # Create large payload (1MB) large_data = "x" * (1024 * 1024) message = {"type": "directive", "data": large_data} # Should handle without error await client._on_directive(message) assert len(client.directives_received) == 1 async def test_concurrent_message_processing(self): """Test handling multiple simultaneous messages""" manager = ConnectionManager() ws = AsyncMock() await manager.connect("client-1", ws) # Send 100 messages concurrently messages = [{"type": "directive", "id": f"cmd-{i}"} for i in range(100)] tasks = [manager.send_to_client("client-1", msg) for msg in messages] await asyncio.gather(*tasks) assert ws.send_json.call_count == 100 async def test_reconnect_backoff_reset(self): """Test that backoff resets after successful connection""" client = BridgeWebSocketClient("ws://test", "key") with patch('websockets.connect', new_callable=AsyncMock) as mock_connect: # First call fails, second succeeds mock_connect.side_effect = [Exception("Fail"), AsyncMock()] with patch('asyncio.sleep', new_callable=AsyncMock): await client.connect() # After successful connection, attempts should reset assert client.reconnect_attempts == 0 async def test_message_type_validation(self): """Test handling of unknown message types""" client = BridgeWebSocketClient("ws://test", "key") # Should not crash on unknown type await client._handle_message({"type": "unknown_type", "data": "test"}) await client._handle_message({"invalid": "no type field"}) async def test_websocket_auth_first_message(self): """Test authentication via first message instead of query param""" # This would be implemented in the actual server pass async def test_client_target_specific_routing(self): """Test that messages can be routed to specific clients""" manager = ConnectionManager() ws1 = AsyncMock() ws2 = AsyncMock() await manager.connect("claude-1", ws1, {"role": "claude"}) await manager.connect("claude-2", ws2, {"role": "claude"}) # Send to specific client await manager.send_to_client("claude-1", {"type": "directive"}) ws1.send_json.assert_called_once() ws2.send_json.assert_not_called() # tests/test_fallback_polling.py import pytest import pytest_asyncio from unittest.mock import AsyncMock, MagicMock, patch import asyncio class PollingFallbackClient: def __init__(self, api_url: str, api_key: str, poll_interval: int = 5): self.api_url = api_url self.api_key = api_key self.poll_interval = poll_interval self.polling = False self.last_poll_time = None async def start_polling(self): self.polling = True while self.polling: await self._poll_once() await asyncio.sleep(self.poll_interval) async def _poll_once(self): # HTTP request to polling endpoint import aiohttp headers = {"X-API-Key": self.api_key} async with aiohttp.ClientSession() as session: async with session.get(f"{self.api_url}/poll", headers=headers) as resp: self.last_poll_time = asyncio.get_event_loop().time() if resp.status == 200: return await resp.json() return None def stop_polling(self): self.polling = False @pytest.mark.asyncio class TestFallbackPolling: async def test_polling_initialization(self): client = PollingFallbackClient("http://api.test", "key-123", 10) assert client.api_url == "http://api.test" assert client.api_key == "key-123" assert client.poll_interval == 10 assert client.polling is False async def test_poll_once_success(self): client = PollingFallbackClient("http://api.test", "key-123") mock_response = MagicMock() mock_response.status = 200 mock_response.json = AsyncMock(return_value={"directives": [{"id": "1"}]}) mock_session = MagicMock() mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) mock_session.get = MagicMock(return_value=mock_response) mock_response.__aenter__ = AsyncMock(return_value=mock_response) mock_response.__aexit__ = AsyncMock(return_value=False) with patch('aiohttp.ClientSession', return_value=mock_session): result = await client._poll_once() assert result == {"directives": [{"id": "1"}]} async def test_poll_once_failure(self): client = PollingFallbackClient("http://api.test", "key-123") mock_response = MagicMock() mock_response.status = 500 mock_session = MagicMock() mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) mock_session.get = MagicMock(return_value=mock_response) mock_response.__aenter__ = AsyncMock(return_value=mock_response) mock_response.__aexit__ = AsyncMock(return_value=False) with patch('aiohttp.ClientSession', return_value=mock_session): result = await client._poll_once() assert result is None async def test_polling_loop(self): client = PollingFallbackClient("http://api.test", "key-123", poll_interval=0.1) call_count = 0 async def mock_poll(): nonlocal call_count call_count += 1 if call_count >= 3: client.stop_polling() return {"id": call_count} with patch.object(client, '_poll_once', side_effect=mock_poll): with patch('asyncio.sleep', new_callable=AsyncMock): await client.start_polling() assert call_count == 3 async def test_polling_stops_on_request(self): client = PollingFallbackClient("http://api.test", "key-123") client.polling = True async def slow_poll(): await asyncio.sleep(10) return {} with patch.object(client, '_poll_once', side_effect=slow_poll): # Stop polling immediately client.stop_polling() assert client.polling is False # tests/test_message_formats.py import pytest import json from pydantic import BaseModel, ValidationError from typing import Optional, Literal from datetime import datetime # Define message schemas for validation class BaseMessage(BaseModel): type: str timestamp: Optional[str] = None class DirectiveMessage(BaseModel): type: Literal["directive"] id: str action: str target: Optional[str] = None parameters: Optional[dict] = None priority: Optional[str] = "normal" class HeartbeatMessage(BaseModel): type: Literal["heartbeat"] timestamp: str class AckMessage(BaseModel): type: Literal["ack"] id: str status: Optional[str] = "received" class StatusMessage(BaseModel): type: Literal["status"] status: str message: Optional[str] = None @pytest.mark.asyncio class TestMessageFormats: def test_directive_message_validation(self): valid = { "type": "directive", "id": "cmd-123", "action": "deploy", "target": "production", "parameters": {"version": "1.0"}, "priority": "high" } msg = DirectiveMessage(**valid) assert msg.id == "cmd-123" assert msg.action == "deploy" assert msg.priority == "high" def test_directive_missing_required_fields(self): invalid = { "type": "directive", "action": "deploy" # missing id } with pytest.raises(ValidationError): DirectiveMessage(**invalid) def test_heartbeat_message(self): msg = HeartbeatMessage( type="heartbeat", timestamp="2024-01-15T10:30:00Z" ) assert msg.type == "heartbeat" def test_ack_message(self): msg = AckMessage( type="ack", id="cmd-123", status="processed" ) assert msg.id == "cmd-123" assert msg.status == "processed" def test_status_message(self): msg = StatusMessage( type="status", status="connected", message="Client ready" ) assert msg.status == "connected" def test_invalid_message_type(self): with pytest.raises(ValidationError): DirectiveMessage( type="invalid_type", # Not "directive" id="cmd-123", action="test" ) def test_malformed_json_handling(self): malformed = "{invalid json" with pytest.raises(json.JSONDecodeError): json.loads(malformed) def test_message_serialization(self): msg = DirectiveMessage( type="directive", id="cmd-123", action="deploy", parameters={"key": "value"} ) json_str = msg.json() parsed = json.loads(json_str) assert parsed["type"] == "directive" assert parsed["id"] == "cmd-123" assert parsed["parameters"]["key"] == "value" # tests/test_end_to_end.py import pytest import pytest_asyncio import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch, Mock, call import asyncpg @pytest.mark.asyncio class TestEndToEnd: @pytest_asyncio.fixture async def full_system(self): """Setup full system with mocked components""" # Mock PostgreSQL mock_pool = MagicMock() mock_conn = AsyncMock() mock_pool.acquire.return_value = mock_conn # Mock WebSocket manager manager = ConnectionManager() return { "db_pool": mock_pool, "conn": mock_conn, "manager": manager } async def test_directive_flow_postgres_to_websocket(self, full_system): """Test complete flow: DB INSERT -> NOTIFY -> WebSocket push""" manager = full_system["manager"] # Setup WebSocket clients ws1 = AsyncMock() ws2 = AsyncMock() await manager.connect("claude-1", ws1) await manager.connect("claude-2", ws2) # Simulate PostgreSQL notification directive = { "type": "directive", "id": "cmd-789", "action": "restart", "target": "service-a" } # Broadcast to all connected clients await manager.broadcast(directive) # Verify both clients received it ws1.send_json.assert_called_with(directive) ws2.send_json.assert_called_with(directive) async def test_client_acknowledgment_flow(self, full_system): """Test client acknowledgment back to server""" manager = full_system["manager"] ws = AsyncMock() await manager.connect("claude-1", ws) # Simulate receiving ack from client ack_message = {"type": "ack", "id": "cmd-123", "status": "processed"} # In real implementation, this would update DB # Here we just verify the flow # Send response back to confirm await manager.send_to_client("claude-1", {"type": "ack_received"}) ws.send_json.assert_called() async def test_multiple_directives_sequencing(self, full_system): """Test that directives are processed in order""" manager = full_system["manager"] ws = AsyncMock() await manager.connect("claude-1", ws) directives = [ {"type": "directive", "id": f"cmd-{i}", "seq": i} for i in range(10) ] for d in directives: await manager.broadcast(d) # Verify order preserved calls = ws.send_json.call_args_list for i, call_args in enumerate(calls): assert call_args[0][0]["seq"] == i async def test_reconnection_during_active_session(self, full_system): """Test client reconnects and resumes receiving""" client = BridgeWebSocketClient("ws://test", "key") with patch('websockets.connect', new_callable=AsyncMock) as mock_ws: mock_socket = AsyncMock() mock_ws.return_value = mock_socket # First connection await client.connect() assert client.connected is True # Simulate disconnect client.connected = False # Reconnect await client.connect() assert client.connected is True async def test_database_trigger_integration(self, full_system): """Test PostgreSQL trigger fires on INSERT""" conn = full_system["conn"] # Simulate the trigger function async def mock_insert_with_trigger(table, data): # Insert data await conn.execute(f"INSERT INTO {table} VALUES ($1)", json.dumps(data)) # Trigger NOTIFY await conn.execute("NOTIFY command_queue, $1", json.dumps(data)) directive = {"id": "cmd-999", "action": "test"} await mock_insert_with_trigger("command_queue", directive) # Verify NOTIFY was called calls = [c for c in conn.execute.call_args_list if 'NOTIFY' in str(c)] assert len(calls) > 0 # tests/test_performance.py import pytest import pytest_asyncio import asyncio import time from unittest.mock import AsyncMock, MagicMock, patch @pytest.mark.asyncio class TestPerformance: async def test_high_frequency_messages(self): """Test handling of high message frequency""" manager = ConnectionManager() # Create 100 clients clients = [AsyncMock() for _ in range(100)] for i, ws in enumerate(clients): await manager.connect(f"client-{i}", ws) # Broadcast 1000 messages start_time = time.time() for i in range(1000): await manager.broadcast({"type": "directive", "seq": i}) end_time = time.time() # Should complete in reasonable time (< 5 seconds for 1000 msgs to 100 clients) assert (end_time - start_time) < 5.0 # Verify all clients received all messages for ws in clients: assert ws.send_json.call_count == 1000 async def test_memory_leak_connections(self): """Test that connections are properly cleaned up""" manager = ConnectionManager() # Rapid connect/disconnect for i in range(1000): ws = AsyncMock() await manager.connect(f"client-{i}", ws) manager.disconnect(f"client-{i}") # Memory should be clean assert manager.get_client_count() == 0 assert len(manager.active_connections) == 0 assert len(manager.client_metadata) == 0 async def test_concurrent_connections(self): """Test handling many concurrent connections""" manager = ConnectionManager() async def connect_client(i): ws = AsyncMock() await manager.connect(f"client-{i}", ws) return ws # Connect 500 clients concurrently tasks = [connect_client(i) for i in range(500)] await asyncio.gather(*tasks) assert manager.get_client_count() == 500 async def test_large_payload_performance(self): """Test performance with large payloads""" manager = ConnectionManager() ws = AsyncMock() await manager.connect("client-1", ws) # 10MB payload large_payload = {"data": "x" * (10 * 1024 * 1024)} start_time = time.time() await manager.broadcast(large_payload) end_time = time.time() # Should handle large payload in reasonable time assert (end_time - start_time) < 2.0 ws.send_json.assert_called_once() # tests/test_security.py import pytest import pytest_asyncio from unittest.mock import AsyncMock, MagicMock, patch import jwt import hashlib @pytest.mark.asyncio class TestSecurity: async def test_token_authentication(self): """Test that valid tokens are accepted""" manager = ConnectionManager() server = BridgeWebSocketServer(manager, "secret-key") ws = AsyncMock() await server.handle_connection(ws, token="secret-key") # Should not close connection assert ws.close.call_count == 0 async def test_invalid_token_rejection(self): """Test that invalid tokens are rejected""" manager = ConnectionManager() server = BridgeWebSocketServer(manager, "secret-key") ws = AsyncMock() await server.handle_connection(ws, token="wrong-key") # Should close connection with policy violation ws.close.assert_called_once() assert ws.close.call_args[1]["code"] == 1008 async def test_sql_injection_protection(self): """Test that channel names are sanitized""" listener = PostgresNotifyListener("postgresql://user:pass@localhost/db") # Attempt injection malicious_channel = "command_queue; DROP TABLE users; --" # Should not execute malicious SQL (this is a simplified check) # In real implementation, use parameterized queries or validation assert ";" in malicious_channel # Detect potential injection async def test_message_forgery_protection(self): """Test that messages are validated before processing""" client = BridgeWebSocketClient("ws://test", "key") # Invalid message structure invalid_messages = [ {"type": "directive"}, # Missing required fields {"id": "123"}, # Missing type {"type": "directive", "id": None}, # Null id {"type": "directive", "id": 123}, # Wrong type for id ] for msg in invalid_messages: # Should handle gracefully without crashing try: await client._handle_message(msg) except Exception as e: # Expected for some invalid formats pass async def test_rate_limiting(self): """Test rate limiting on connections""" manager = ConnectionManager() # Rapid connections from same source connections = [] for i in range(100): ws = AsyncMock() await manager.connect(f"client-{i}", ws) connections.append(ws) # All should be accepted (rate limiting would be at higher level) assert manager.get_client_count() == 100 async def test_secure_header_handling(self): """Test that API keys are handled securely in headers""" # API key should not be logged api_key = "secret-key-123" # Simulate header processing headers = {"X-API-Key": api_key} # Key should be present but not exposed in logs assert headers["X-API-Key"] == api_key # In real implementation, would verify no logging of key # tests/test_cli_script.py import pytest import subprocess import sys from unittest.mock import patch, MagicMock import argparse # Test the CLI script that would run the WebSocket client class TestCLI: def test_cli_argument_parsing(self): """Test CLI argument parsing""" test_args = [ "--uri", "ws://localhost:8000/bridge/ws", "--api-key", "test-key", "--fallback-poll", "10" ] # Parse args parser = argparse.ArgumentParser() parser.add_argument("--uri", default="ws://localhost:8000/bridge/ws") parser.add_argument("--api-key", required=True) parser.add_argument("--fallback-poll", type=int, default=5) args = parser.parse_args(test_args) assert args.uri == "ws://localhost:8000/bridge/ws" assert args.api_key == "test-key" assert args.fallback_poll == 10 def test_cli_missing_api_key(self): """Test CLI fails without API key""" parser = argparse.ArgumentParser() parser.add_argument("--api-key", required=True) with pytest.raises(SystemExit): parser.parse_args([]) @patch('asyncio.run') def test_cli_main_execution(self, mock_run): """Test main execution path""" # This would test the actual main() function pass # Run configuration for pytest def pytest_configure(config): config.addinivalue_line("markers", "slow: marks tests as slow") config.addinivalue_line("markers", "integration: marks tests as integration tests")