test_client.py

#!/usr/bin/env python3
"""
Test client for STT GPU Service on HuggingFace Spaces
Tests the transcription API with audio files or generates test audio
"""

import os
import sys
import base64
import json
import requests
import tempfile
import time
from typing import Optional

# HuggingFace Space URL
STT_SERVICE_URL = "https://pgits-stt-gpu-service-v2.hf.space"

def test_health_check():
    """Test the health endpoint"""
    print("🔍 Testing health check...")
    try:
        response = requests.get(f"{STT_SERVICE_URL}/api/health", timeout=10)
        if response.status_code == 200:
            print("✅ Health check passed")
            return True
        else:
            print(f"❌ Health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Health check error: {e}")
        return False

def create_test_audio_file() -> str:
    """Create a simple test audio file using system tools"""
    print("🎵 Creating test audio file...")
    
    try:
        # Try to create a simple sine wave audio file using ffmpeg if available
        test_audio_path = "/tmp/test_audio.wav"
        
        # Generate 3 seconds of 440Hz sine wave
        os.system(f"ffmpeg -f lavfi -i 'sine=frequency=440:duration=3' -y {test_audio_path} 2>/dev/null")
        
        if os.path.exists(test_audio_path) and os.path.getsize(test_audio_path) > 1000:
            print(f"✅ Test audio created: {test_audio_path}")
            return test_audio_path
        else:
            print("⚠️  ffmpeg not available or failed, using minimal test data")
            return None
            
    except Exception as e:
        print(f"⚠️  Could not create test audio: {e}")
        return None

def test_transcription_with_file(audio_file_path: str, language: str = "en", model_size: str = "base"):
    """Test transcription with an audio file"""
    print(f"🎤 Testing transcription with file: {audio_file_path}")
    
    try:
        # Read and encode audio file
        with open(audio_file_path, "rb") as f:
            audio_data = f.read()
            audio_base64 = base64.b64encode(audio_data).decode('utf-8')
        
        print(f"📄 Audio file size: {len(audio_data)} bytes")
        print(f"📄 Base64 size: {len(audio_base64)} characters")
        
        # Send transcription request
        payload = {
            "audio_base64": audio_base64,
            "language": language,
            "model_size": model_size
        }
        
        print(f"🚀 Sending request to {STT_SERVICE_URL}/api/transcribe")
        start_time = time.time()
        
        response = requests.post(
            f"{STT_SERVICE_URL}/api/transcribe",
            json=payload,
            timeout=30,
            headers={"Content-Type": "application/json"}
        )
        
        processing_time = time.time() - start_time
        
        print(f"📊 Request completed in {processing_time:.2f}s")
        print(f"📊 Response status: {response.status_code}")
        
        if response.status_code == 200:
            try:
                result = response.json()
                print(f"✅ Transcription result: {result}")
                return result
            except json.JSONDecodeError:
                print(f"✅ Transcription result (text): {response.text}")
                return response.text
        else:
            print(f"❌ Request failed: {response.status_code}")
            print(f"❌ Error response: {response.text}")
            return None
            
    except Exception as e:
        print(f"❌ Transcription test failed: {e}")
        return None

def test_transcription_with_minimal_data():
    """Test transcription with minimal test data"""
    print("🧪 Testing with minimal base64 data...")
    
    # Create minimal test data (should trigger demo response)
    test_data = b"test audio data for demo"
    audio_base64 = base64.b64encode(test_data).decode('utf-8')
    
    payload = {
        "audio_base64": audio_base64,
        "language": "en",
        "model_size": "base"
    }
    
    try:
        response = requests.post(
            f"{STT_SERVICE_URL}/api/transcribe",
            json=payload,
            timeout=10,
            headers={"Content-Type": "application/json"}
        )
        
        if response.status_code == 200:
            try:
                result = response.json()
                print(f"✅ Demo test result: {result}")
            except json.JSONDecodeError:
                print(f"✅ Demo test result (text): {response.text}")
        else:
            print(f"❌ Demo test failed: {response.status_code} - {response.text}")
            
    except Exception as e:
        print(f"❌ Demo test error: {e}")

def main():
    """Main test function"""
    print("🎤 STT GPU Service Test Client")
    print("=" * 50)
    
    # Test health check first
    if not test_health_check():
        print("❌ Service appears to be down, continuing with other tests...")
    
    print()
    
    # Test with minimal data first
    test_transcription_with_minimal_data()
    print()
    
    # Check if user provided an audio file
    if len(sys.argv) > 1:
        audio_file = sys.argv[1]
        if os.path.exists(audio_file):
            print(f"🎵 Using provided audio file: {audio_file}")
            test_transcription_with_file(audio_file)
        else:
            print(f"❌ Audio file not found: {audio_file}")
    else:
        # Try to create test audio
        test_audio = create_test_audio_file()
        if test_audio:
            test_transcription_with_file(test_audio)
            # Clean up
            try:
                os.unlink(test_audio)
            except:
                pass
        else:
            print("⚠️  No audio file provided and couldn't create test audio")
            print("💡 Usage: python test_client.py [audio_file.wav]")
            print("💡 Or install ffmpeg to auto-generate test audio")
    
    print()
    print("🎯 Test complete!")
    print("💡 You can also test via web interface at:")
    print(f"   {STT_SERVICE_URL}")

if __name__ == "__main__":
    main()