test_websocket.py

import asyncio
import websockets
import json
import numpy as np
import time

# Configuration
HF_SPACE_URL = "https://androidguy-speaker-diarization.hf.space"
RENDER_SIGNAL_URL = "wss://render-signal-audio.onrender.com/stream"
WS_INFERENCE_ENDPOINT = f"wss://{HF_SPACE_URL.replace('https://', '')}/ws_inference"
WS_RELAY_ENDPOINT = "wss://render-signal-audio.onrender.com/ws_relay"

async def test_websocket_connections():
    print(f"Testing WebSocket connections...")
    
    # Test HF Space /ws_inference connection
    print(f"Connecting to {WS_INFERENCE_ENDPOINT}...")
    try:
        async with websockets.connect(WS_INFERENCE_ENDPOINT, ping_interval=30) as ws:
            print("Connected to HF Space inference WebSocket!")
            
            # Send test data (simulated audio)
            test_audio = np.zeros(1600, dtype=np.float32).tobytes()
            await ws.send(test_audio)
            print(f"Sent test audio data: {len(test_audio)} bytes")
            
            # Wait for response
            for _ in range(5):  # Try to get 5 responses
                try:
                    response = await asyncio.wait_for(ws.recv(), timeout=3.0)
                    if isinstance(response, str):
                        try:
                            json_data = json.loads(response)
                            print(f"Received JSON response: {json.dumps(json_data, indent=2)}")
                        except json.JSONDecodeError:
                            print(f"Received non-JSON response: {response[:500]}...")
                    else:
                        print(f"Received binary data: {len(response)} bytes")
                except asyncio.TimeoutError:
                    print("No response received within timeout")
                    break
                
    except Exception as e:
        print(f"Error connecting to HF Space WebSocket: {e}")
        
    # Test Render Signal WebSocket relay connection
    print(f"\nConnecting to {WS_RELAY_ENDPOINT}...")
    try:
        async with websockets.connect(WS_RELAY_ENDPOINT) as ws:
            print("Connected to Render Signal WebSocket relay!")
            
            # Send test configuration message
            test_config = {
                "type": "config",
                "client_id": f"test_client_{time.time()}",
                "test": True
            }
            await ws.send(json.dumps(test_config))
            print(f"Sent test config: {test_config}")
            
            # Wait for responses
            for _ in range(5):  # Try to get 5 responses
                try:
                    response = await asyncio.wait_for(ws.recv(), timeout=3.0)
                    if isinstance(response, str):
                        try:
                            json_data = json.loads(response)
                            print(f"Received JSON response: {json.dumps(json_data, indent=2)}")
                        except json.JSONDecodeError:
                            print(f"Received non-JSON response: {response[:500]}...")
                    else:
                        print(f"Received binary data: {len(response)} bytes")
                except asyncio.TimeoutError:
                    print("No response received within timeout")
                    break
    except Exception as e:
        print(f"Error connecting to Render Signal WebSocket: {e}")

# Run tests for direct WebSocket communication
if __name__ == "__main__":
    asyncio.run(test_websocket_connections())