import pytest import asyncio import json import websockets from src.api_hooks import WebSocketServer @pytest.mark.asyncio async def test_websocket_subscription_and_broadcast(): # Mock app app = type("MockApp", (), {"test_hooks_enabled": True})() # Start server on a specific port port = 9005 server = WebSocketServer(app, port=port) server.start() # Wait for server to start await asyncio.sleep(0.5) try: uri = f"ws://127.0.0.1:{port}" async with websockets.connect(uri) as websocket: # Subscribe to events channel subscribe_msg = {"action": "subscribe", "channel": "events"} await websocket.send(json.dumps(subscribe_msg)) # Receive confirmation response = await websocket.recv() data = json.loads(response) assert data["type"] == "subscription_confirmed" assert data["channel"] == "events" # Broadcast an event from the server event_payload = {"event": "test_event", "data": "hello"} server.broadcast("events", event_payload) # Receive the broadcast broadcast_response = await websocket.recv() broadcast_data = json.loads(broadcast_response) assert broadcast_data["channel"] == "events" assert broadcast_data["payload"] == event_payload finally: server.stop()