#!/usr/bin/env python3 """ WebSocket TTS API Test Script Tests connection, latency, and audio streaming """ import asyncio import websockets import json import time from pathlib import Path # Configuration WS_URL = "wss://ebitlogix-parler-tts-api.hf.space/ws/tts" # For local testing: WS_URL = "ws://localhost:7860/ws/tts" OUTPUT_DIR = Path("tts_output") OUTPUT_DIR.mkdir(exist_ok=True) async def test_websocket_connection(): """Test basic WebSocket connection""" print("\n" + "="*60) print("Testing WebSocket Connection") print("="*60) try: async with websockets.connect(WS_URL) as websocket: print("✅ Connected to WebSocket server") print(f" URL: {WS_URL}") print(f" Connection state: {websocket.state}") return True except Exception as e: print(f"❌ Failed to connect: {e}") return False async def test_websocket_response(): """Test WebSocket request/response""" print("\n" + "="*60) print("Testing WebSocket Request/Response") print("="*60) try: async with websockets.connect(WS_URL) as websocket: # Send request request = { "text": "سلام دنیا", "speaker": "Divya", "pitch": "Moderate", "rate": "Moderate" } print(f"\nSending request:") print(f" Text: {request['text']}") print(f" Speaker: {request['speaker']}") await websocket.send(json.dumps(request)) print("✅ Request sent") # Receive responses responses = [] while True: try: message = await asyncio.wait_for(websocket.recv(), timeout=60) if isinstance(message, str): data = json.loads(message) print(f"\n📨 Status Message: {data}") responses.append(data) if data.get("status") == "complete": print(f"✅ Generation complete! Received {data['chunks_sent']} chunks") break else: # Binary audio data print(f"🔊 Audio chunk received: {len(message)} bytes") responses.append({"type": "audio", "size": len(message)}) except asyncio.TimeoutError: print("❌ Timeout waiting for response") break return len(responses) > 0 except Exception as e: print(f"❌ Error: {e}") import traceback traceback.print_exc() return False async def test_websocket_latency(): """Test WebSocket latency and streaming speed""" print("\n" + "="*60) print("Testing WebSocket Latency & Streaming Speed") print("="*60) try: start_time = time.time() async with websockets.connect(WS_URL) as websocket: connection_time = time.time() - start_time print(f"\n✅ Connection established in {connection_time*1000:.1f}ms") # Send request request = { "text": "یہ ایک ٹیسٹ ہے", "speaker": "Rani" } send_time = time.time() await websocket.send(json.dumps(request)) print(f"✅ Request sent in {(time.time()-send_time)*1000:.1f}ms") # Track first chunk time first_chunk_time = None total_audio_size = 0 chunk_count = 0 while True: message = await asyncio.wait_for(websocket.recv(), timeout=60) if isinstance(message, str): data = json.loads(message) if data.get("status") == "generating": print(f"📊 Status: {data.get('message')}") elif data.get("status") == "complete": total_time = time.time() - start_time print(f"\n✅ Complete!") print(f" Total time: {total_time:.2f}s") print(f" First chunk: {first_chunk_time*1000:.1f}ms") print(f" Total chunks: {chunk_count}") print(f" Total audio size: {total_audio_size/1024:.1f} KB") if chunk_count > 0: print(f" Avg chunk size: {total_audio_size/chunk_count:.0f} bytes") break else: # Audio chunk if first_chunk_time is None: first_chunk_time = time.time() - start_time print(f"\n🔊 First audio chunk received in {first_chunk_time*1000:.1f}ms") total_audio_size += len(message) chunk_count += 1 print(f" Chunk {chunk_count}: {len(message)} bytes") return True except Exception as e: print(f"❌ Error: {e}") import traceback traceback.print_exc() return False async def test_websocket_streaming_save(): """Test WebSocket streaming and save audio""" print("\n" + "="*60) print("Testing WebSocket Streaming & Audio Save") print("="*60) try: async with websockets.connect(WS_URL) as websocket: request = { "text": "مرحبا، یہ ایک WebSocket ٹیسٹ ہے", "speaker": "Generic Female" } print(f"\nSending: {request['text']}") await websocket.send(json.dumps(request)) # Collect all audio chunks audio_chunks = [] chunk_count = 0 while True: message = await asyncio.wait_for(websocket.recv(), timeout=60) if isinstance(message, str): data = json.loads(message) print(f"Status: {data}") if data.get("status") == "complete": break else: # Audio chunk audio_chunks.append(message) chunk_count += 1 print(f"Received chunk {chunk_count}: {len(message)} bytes") # Save combined audio if audio_chunks: combined_audio = b"".join(audio_chunks) timestamp = time.strftime("%Y%m%d_%H%M%S") filename = OUTPUT_DIR / f"websocket_test_{timestamp}.wav" with open(filename, "wb") as f: f.write(combined_audio) print(f"\n✅ Audio saved: {filename}") print(f" Total size: {len(combined_audio)/1024:.1f} KB") return True except Exception as e: print(f"❌ Error: {e}") import traceback traceback.print_exc() return False async def main(): """Run all tests""" print("\n" + "█"*60) print("█ WebSocket TTS API Test Suite") print("█"*60) # Run tests tests = [ ("Connection Test", test_websocket_connection()), ("Request/Response Test", test_websocket_response()), ("Latency Test", test_websocket_latency()), ("Streaming & Save Test", test_websocket_streaming_save()), ] results = {} for test_name, test_coro in tests: try: results[test_name] = await test_coro except Exception as e: print(f"\n❌ {test_name} failed: {e}") results[test_name] = False # Summary print("\n" + "="*60) print("Test Summary") print("="*60) for test_name, result in results.items(): status = "✅ PASSED" if result else "❌ FAILED" print(f"{status} - {test_name}") passed = sum(1 for r in results.values() if r) print(f"\nTotal: {passed}/{len(results)} tests passed") print("="*60 + "\n") if __name__ == "__main__": asyncio.run(main())