#!/usr/bin/env python3 """ Local testing script for omega_model server. Tests the Docker container health and API endpoints. """ import requests import time import numpy as np import base64 import io import sys from pathlib import Path # Configuration SERVER_URL = "http://localhost:8000" HEALTH_ENDPOINT = f"{SERVER_URL}/api/v1/health" V2V_ENDPOINT = f"{SERVER_URL}/api/v1/v2v" V2T_ENDPOINT = f"{SERVER_URL}/api/v1/v2t" # Wait times MAX_WAIT_TIME = 300 # 5 minutes for model loading HEALTH_CHECK_INTERVAL = 5 # Check every 5 seconds def create_test_audio(duration_seconds=2, sample_rate=16000): """Create a simple test audio signal (sine wave).""" t = np.linspace(0, duration_seconds, int(sample_rate * duration_seconds)) # Generate a 440 Hz sine wave (A note) audio = np.sin(2 * np.pi * 440 * t).astype(np.float32) # Reshape to (1, samples) format audio = audio.reshape(1, -1) return audio, sample_rate def audio_to_base64(audio, sample_rate): """Convert numpy audio array to base64 string.""" buf = io.BytesIO() np.save(buf, audio.astype(np.float32)) return base64.b64encode(buf.getvalue()).decode() def wait_for_server(max_wait=MAX_WAIT_TIME): """Wait for the server to be ready.""" print(f"Waiting for server to be ready (max {max_wait}s)...") start_time = time.time() while time.time() - start_time < max_wait: try: response = requests.get(HEALTH_ENDPOINT, timeout=5) if response.status_code == 200: data = response.json() if data.get("status") == "healthy" and data.get("model_loaded"): print("✅ Server is ready!") return True else: print(f"⏳ Server starting... (status: {data.get('status')}, model_loaded: {data.get('model_loaded')})") else: print(f"⏳ Server responding but not ready (status code: {response.status_code})") except requests.exceptions.ConnectionError: print("⏳ Server not responding yet...") except Exception as e: print(f"⏳ Error checking server: {e}") time.sleep(HEALTH_CHECK_INTERVAL) print("❌ Server did not become ready in time!") return False def test_health(): """Test the health endpoint.""" print("\n" + "="*60) print("TEST 1: Health Check") print("="*60) try: response = requests.get(HEALTH_ENDPOINT, timeout=10) response.raise_for_status() data = response.json() print(f"Status Code: {response.status_code}") print(f"Response: {data}") if data.get("status") == "healthy" and data.get("model_loaded"): print("✅ Health check PASSED") return True else: print("❌ Health check FAILED") if data.get("error"): print(f"Error: {data.get('error')}") return False except Exception as e: print(f"❌ Health check FAILED: {e}") return False def test_v2t(): """Test the voice-to-text endpoint.""" print("\n" + "="*60) print("TEST 2: Voice-to-Text (v2t)") print("="*60) try: # Create test audio audio, sample_rate = create_test_audio(duration_seconds=2) audio_b64 = audio_to_base64(audio, sample_rate) payload = { "audio_data": audio_b64, "sample_rate": sample_rate } print(f"Sending request with audio shape: {audio.shape}, sample_rate: {sample_rate}") print("Waiting for response (this may take 30-60 seconds)...") response = requests.post(V2T_ENDPOINT, json=payload, timeout=120) response.raise_for_status() data = response.json() print(f"Status Code: {response.status_code}") print(f"Transcribed Text: {data.get('text', 'N/A')}") if "text" in data: print("✅ Voice-to-Text test PASSED") return True else: print("❌ Voice-to-Text test FAILED - no text in response") return False except requests.exceptions.Timeout: print("❌ Voice-to-Text test FAILED - Request timed out") return False except Exception as e: print(f"❌ Voice-to-Text test FAILED: {e}") import traceback traceback.print_exc() return False def test_v2v(): """Test the voice-to-voice endpoint.""" print("\n" + "="*60) print("TEST 3: Voice-to-Voice (v2v)") print("="*60) try: # Create test audio audio, sample_rate = create_test_audio(duration_seconds=2) audio_b64 = audio_to_base64(audio, sample_rate) payload = { "audio_data": audio_b64, "sample_rate": sample_rate } print(f"Sending request with audio shape: {audio.shape}, sample_rate: {sample_rate}") print("Waiting for response (this may take 60-120 seconds)...") response = requests.post(V2V_ENDPOINT, json=payload, timeout=200) response.raise_for_status() data = response.json() print(f"Status Code: {response.status_code}") if "audio_data" in data: # Decode the response audio audio_response_b64 = data["audio_data"] audio_response_bytes = base64.b64decode(audio_response_b64) audio_response = np.load(io.BytesIO(audio_response_bytes), allow_pickle=False) print(f"Response audio shape: {audio_response.shape}") print(f"Response audio dtype: {audio_response.dtype}") print("✅ Voice-to-Voice test PASSED") return True else: print("❌ Voice-to-Voice test FAILED - no audio_data in response") return False except requests.exceptions.Timeout: print("❌ Voice-to-Voice test FAILED - Request timed out") return False except Exception as e: print(f"❌ Voice-to-Voice test FAILED: {e}") import traceback traceback.print_exc() return False def main(): """Run all tests.""" print("="*60) print("OMEGA MODEL LOCAL TESTING") print("="*60) print(f"Server URL: {SERVER_URL}") print(f"Make sure the Docker container is running on port 8010") print() # Wait for server to be ready if not wait_for_server(): print("\n❌ Server is not ready. Please check Docker container logs.") sys.exit(1) # Run tests results = [] results.append(("Health Check", test_health())) results.append(("Voice-to-Text", test_v2t())) results.append(("Voice-to-Voice", test_v2v())) # Summary print("\n" + "="*60) print("TEST SUMMARY") print("="*60) passed = sum(1 for _, result in results if result) total = len(results) for test_name, result in results: status = "✅ PASSED" if result else "❌ FAILED" print(f"{test_name}: {status}") print(f"\nTotal: {passed}/{total} tests passed") if passed == total: print("\n🎉 All tests passed! Your model is ready to upload.") sys.exit(0) else: print("\n⚠️ Some tests failed. Please check the errors above.") sys.exit(1) if __name__ == "__main__": main()