test_websocket.py

#!/usr/bin/env python3
"""
Test script to verify WebSocket real-time functionality
"""

import asyncio
import websockets
import json
import sys

async def test_websocket():
    """Test the WebSocket connection for real-time updates"""
    uri = "ws://127.0.0.1:8000/ws/updates"

    try:
        async with websockets.connect(uri) as websocket:
            print("✅ WebSocket connected successfully")

            # Send a ping message
            await websocket.send(json.dumps({"type": "ping"}))
            response = await websocket.recv()
            data = json.loads(response)
            print(f"✅ Ping/pong test: {data.get('event', 'unknown')}")

            # Listen for real-time updates for 10 seconds
            print("🔄 Listening for real-time updates...")
            start_time = asyncio.get_event_loop().time()

            while asyncio.get_event_loop().time() - start_time < 10:
                try:
                    message = await asyncio.wait_for(websocket.recv(), timeout=1.0)
                    data = json.loads(message)
                    print(f"📡 Received: {data.get('event', 'unknown')} - {data.get('type', 'unknown')}")
                except asyncio.TimeoutError:
                    pass  # No message received, continue listening

            print("✅ WebSocket test completed successfully")

    except Exception as e:
        print(f"❌ WebSocket test failed: {e}")
        return False

    return True

if __name__ == "__main__":
    success = asyncio.run(test_websocket())
    if success:
        print("\n🎉 Real-time WebSocket functionality is working!")
    else:
        print("\n⚠️ WebSocket functionality needs fixes")
        sys.exit(1)