test_websocket.py

#!/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())