test_local.py

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