|
| 1 | +#!/usr/bin/env python |
| 2 | +""" |
| 3 | +Test WebSocket connection to TTSFM server. |
| 4 | +
|
| 5 | +This script tests the WebSocket functionality by connecting to the server |
| 6 | +and performing a simple TTS generation request. |
| 7 | +""" |
| 8 | + |
| 9 | +import time |
| 10 | +import socketio |
| 11 | + |
| 12 | +# Create a Socket.IO client |
| 13 | +sio = socketio.Client(logger=True, engineio_logger=True) |
| 14 | + |
| 15 | +# Track connection state |
| 16 | +connected = False |
| 17 | +stream_complete = False |
| 18 | +chunks_received = 0 |
| 19 | + |
| 20 | + |
| 21 | +@sio.on('connect') |
| 22 | +def on_connect(): |
| 23 | + """Handle connection event.""" |
| 24 | + global connected |
| 25 | + connected = True |
| 26 | + print('\n✅ Connected to WebSocket server!') |
| 27 | + print(f'Session ID: {sio.sid}') |
| 28 | + |
| 29 | + |
| 30 | +@sio.on('connected') |
| 31 | +def on_session_ready(data): |
| 32 | + """Handle session ready event.""" |
| 33 | + print(f'\n✅ Session established: {data}') |
| 34 | + |
| 35 | + |
| 36 | +@sio.on('disconnect') |
| 37 | +def on_disconnect(): |
| 38 | + """Handle disconnection event.""" |
| 39 | + global connected |
| 40 | + connected = False |
| 41 | + print('\n❌ Disconnected from WebSocket server') |
| 42 | + |
| 43 | + |
| 44 | +@sio.on('connect_error') |
| 45 | +def on_connect_error(data): |
| 46 | + """Handle connection error.""" |
| 47 | + print(f'\n❌ Connection error: {data}') |
| 48 | + |
| 49 | + |
| 50 | +@sio.on('pong') |
| 51 | +def on_pong(data): |
| 52 | + """Handle pong response.""" |
| 53 | + print(f'\n✅ Pong received: {data}') |
| 54 | + |
| 55 | + |
| 56 | +@sio.on('stream_started') |
| 57 | +def on_stream_started(data): |
| 58 | + """Handle stream started event.""" |
| 59 | + print(f'\n✅ Stream started: {data}') |
| 60 | + |
| 61 | + |
| 62 | +@sio.on('stream_progress') |
| 63 | +def on_stream_progress(data): |
| 64 | + """Handle stream progress event.""" |
| 65 | + progress = data.get('progress', 0) |
| 66 | + status = data.get('status', 'unknown') |
| 67 | + print(f'📊 Progress: {progress}% - Status: {status}') |
| 68 | + |
| 69 | + |
| 70 | +@sio.on('audio_chunk') |
| 71 | +def on_audio_chunk(data): |
| 72 | + """Handle audio chunk event.""" |
| 73 | + global chunks_received |
| 74 | + chunks_received += 1 |
| 75 | + chunk_index = data.get('chunk_index', 0) |
| 76 | + total_chunks = data.get('total_chunks', 0) |
| 77 | + print(f'🎵 Received audio chunk {chunk_index + 1}/{total_chunks}') |
| 78 | + |
| 79 | + |
| 80 | +@sio.on('stream_complete') |
| 81 | +def on_stream_complete(data): |
| 82 | + """Handle stream complete event.""" |
| 83 | + global stream_complete |
| 84 | + stream_complete = True |
| 85 | + print(f'\n✅ Stream complete: {data}') |
| 86 | + print(f'Total chunks received: {chunks_received}') |
| 87 | + |
| 88 | + |
| 89 | +@sio.on('stream_error') |
| 90 | +def on_stream_error(data): |
| 91 | + """Handle stream error event.""" |
| 92 | + print(f'\n❌ Stream error: {data}') |
| 93 | + |
| 94 | + |
| 95 | +def test_connection(url='http://localhost:8000'): |
| 96 | + """Test WebSocket connection.""" |
| 97 | + print(f'🔌 Connecting to {url}...') |
| 98 | + |
| 99 | + try: |
| 100 | + # Connect to the server |
| 101 | + sio.connect(url, transports=['polling', 'websocket']) |
| 102 | + |
| 103 | + # Wait for connection |
| 104 | + timeout = 10 |
| 105 | + start_time = time.time() |
| 106 | + while not connected and (time.time() - start_time) < timeout: |
| 107 | + time.sleep(0.1) |
| 108 | + |
| 109 | + if not connected: |
| 110 | + print('❌ Failed to connect within timeout') |
| 111 | + return False |
| 112 | + |
| 113 | + # Test ping/pong |
| 114 | + print('\n📡 Testing ping/pong...') |
| 115 | + sio.emit('ping', {'timestamp': time.time()}) |
| 116 | + time.sleep(1) |
| 117 | + |
| 118 | + # Test TTS generation |
| 119 | + print('\n🎤 Testing TTS generation...') |
| 120 | + request_data = { |
| 121 | + 'request_id': f'test_{int(time.time())}', |
| 122 | + 'text': 'Hello, this is a WebSocket test!', |
| 123 | + 'voice': 'alloy', |
| 124 | + 'format': 'mp3', |
| 125 | + 'chunk_size': 512 |
| 126 | + } |
| 127 | + |
| 128 | + sio.emit('generate_stream', request_data) |
| 129 | + |
| 130 | + # Wait for stream to complete |
| 131 | + timeout = 30 |
| 132 | + start_time = time.time() |
| 133 | + while not stream_complete and (time.time() - start_time) < timeout: |
| 134 | + time.sleep(0.1) |
| 135 | + |
| 136 | + if stream_complete: |
| 137 | + print('\n✅ WebSocket test completed successfully!') |
| 138 | + return True |
| 139 | + else: |
| 140 | + print('\n⚠️ Stream did not complete within timeout') |
| 141 | + return False |
| 142 | + |
| 143 | + except Exception as e: |
| 144 | + print(f'\n❌ Error during test: {e}') |
| 145 | + import traceback |
| 146 | + traceback.print_exc() |
| 147 | + return False |
| 148 | + |
| 149 | + finally: |
| 150 | + # Disconnect |
| 151 | + if connected: |
| 152 | + print('\n🔌 Disconnecting...') |
| 153 | + sio.disconnect() |
| 154 | + time.sleep(1) |
| 155 | + |
| 156 | + |
| 157 | +if __name__ == '__main__': |
| 158 | + import sys |
| 159 | + |
| 160 | + # Get URL from command line or use default |
| 161 | + url = sys.argv[1] if len(sys.argv) > 1 else 'http://localhost:8000' |
| 162 | + |
| 163 | + print('=' * 60) |
| 164 | + print('TTSFM WebSocket Connection Test') |
| 165 | + print('=' * 60) |
| 166 | + |
| 167 | + success = test_connection(url) |
| 168 | + |
| 169 | + print('\n' + '=' * 60) |
| 170 | + if success: |
| 171 | + print('✅ All tests passed!') |
| 172 | + sys.exit(0) |
| 173 | + else: |
| 174 | + print('❌ Some tests failed') |
| 175 | + sys.exit(1) |
| 176 | + |
0 commit comments