test_endpoints.py

"""
Test file for Multi-Utility Server endpoints.
Test both Subtitles and Embeddings APIs with the deployed server.

Usage:
    python test_endpoints.py

Modify the test data in the TEST_DATA section to test different inputs.
"""

import json
import os
from pathlib import Path
from typing import Any, Dict, List

import requests

# ============================================================================
# CONFIGURATION
# ============================================================================

BASE_URL = "https://abhisheksan-multiutility-server.hf.space"
API_KEY = "zL6Rod5869J0"  # Replace with your actual API key

# ============================================================================
# TEST DATA - MODIFY THESE TO TEST DIFFERENT INPUTS
# ============================================================================

# Test data for Subtitles API
SUBTITLE_TEST_DATA = {
    "url": "https://www.youtube.com/watch?v=dQw4w9WgXcQ",
    "lang": "en",
}

# Test data for Embeddings API
EMBEDDING_TEST_DATA = {
    "texts": [
        "Hello, how are you?",
        "Machine learning is fascinating",
        "Python is a great programming language",
    ],
    "normalize": True,
}

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================


def print_section(title: str):
    """Print a formatted section header."""
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80 + "\n")


def print_response(response: requests.Response):
    """Print formatted response details."""
    print(f"Status Code: {response.status_code}")
    print(f"Response Time: {response.elapsed.total_seconds():.2f}s")
    print("\nResponse Headers:")
    for key, value in response.headers.items():
        print(f"  {key}: {value}")
    print("\nResponse Body:")
    try:
        print(json.dumps(response.json(), indent=2))
    except Exception:
        print(response.text)


def make_request(
    method: str,
    endpoint: str,
    data: Dict[str, Any] = None,
    headers: Dict[str, str] = None,
) -> requests.Response:
    """Make HTTP request to the API."""
    url = f"{BASE_URL}{endpoint}"
    default_headers = {"Content-Type": "application/json", "X-API-Key": API_KEY}

    if headers:
        default_headers.update(headers)

    try:
        if method.upper() == "GET":
            response = requests.get(url, headers=default_headers)
        elif method.upper() == "POST":
            response = requests.post(url, json=data, headers=default_headers)
        else:
            raise ValueError(f"Unsupported HTTP method: {method}")

        return response
    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {str(e)}")
        raise


# ============================================================================
# TEST FUNCTIONS
# ============================================================================


def test_health_check():
    """Test the main health check endpoint."""
    print_section("TEST: Health Check")

    try:
        response = make_request("GET", "/health")
        print_response(response)

        if response.status_code == 200:
            print("\n✅ Health check passed!")
        else:
            print("\n❌ Health check failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")


def test_root_endpoint():
    """Test the root endpoint."""
    print_section("TEST: Root Endpoint")

    try:
        response = make_request("GET", "/")
        print_response(response)

        if response.status_code == 200:
            print("\n✅ Root endpoint test passed!")
        else:
            print("\n❌ Root endpoint test failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")


def test_subtitles_health():
    """Test the subtitles service health check."""
    print_section("TEST: Subtitles Health Check")

    try:
        response = make_request("GET", "/api/v1/subtitles/health")
        print_response(response)

        if response.status_code == 200:
            print("\n✅ Subtitles health check passed!")
        else:
            print("\n❌ Subtitles health check failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")


def test_audio_transcription():
    """Test the audio file transcription endpoint."""
    print_section("TEST: Audio File Transcription")

    # Create a small test audio file (silent WAV)
    # This is a minimal WAV file header for a 1-second silent audio
    print("Creating test audio file...")
    test_audio = b"RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00D\xac\x00\x00\x88X\x01\x00\x02\x00\x10\x00data\x00\x00\x00\x00"

    test_file_path = "test_audio.wav"

    try:
        # Write test audio file
        with open(test_file_path, "wb") as f:
            f.write(test_audio)

        print(f"Test file created: {test_file_path}")
        print("Note: Using minimal WAV file for testing\n")

        # Make request with file upload
        url = f"{BASE_URL}/api/v1/subtitles/transcribe"
        headers = {"X-API-Key": API_KEY}

        with open(test_file_path, "rb") as f:
            files = {"file": ("test_audio.wav", f, "audio/wav")}
            data = {"lang": "en"}

            response = requests.post(url, headers=headers, files=files, data=data)

        print_response(response)

        if response.status_code == 200:
            data = response.json()
            transcription_count = len(data.get("transcription", []))
            print(f"\n✅ Audio transcription passed!")
            print(f"   - Language: {data.get('language')}")
            print(f"   - File name: {data.get('file_name')}")
            print(f"   - Number of segments: {transcription_count}")

            # Print transcription segments
            if transcription_count > 0:
                print(f"\n   Transcription segments:")
                for i, line in enumerate(data.get("transcription", [])[:5]):
                    print(f"     {i + 1}. {line}")
            else:
                print("   (No transcription - test audio is silent)")
        else:
            print("\n❌ Audio transcription failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")
    finally:
        # Cleanup test file
        if os.path.exists(test_file_path):
            os.remove(test_file_path)
            print(f"\n   Cleaned up test file: {test_file_path}")


def test_subtitles_extract():
    """Test the subtitles extraction endpoint."""
    print_section("TEST: Extract Subtitles")

    print("Test Data:")
    print(json.dumps(SUBTITLE_TEST_DATA, indent=2))
    print()

    try:
        response = make_request(
            "POST", "/api/v1/subtitles/extract", data=SUBTITLE_TEST_DATA
        )
        print_response(response)

        if response.status_code == 200:
            data = response.json()
            subtitle_count = len(data.get("subtitles", []))
            print(f"\n✅ Subtitles extraction passed!")
            print(f"   - Video ID: {data.get('video_id')}")
            print(f"   - Language: {data.get('language')}")
            print(f"   - Number of subtitle lines: {subtitle_count}")

            # Print first few subtitle lines
            if subtitle_count > 0:
                print(f"\n   First 5 subtitle lines:")
                for i, line in enumerate(data.get("subtitles", [])[:5]):
                    print(f"     {i + 1}. {line}")
        else:
            print("\n❌ Subtitles extraction failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")


def test_embeddings_health():
    """Test the embeddings service health check."""
    print_section("TEST: Embeddings Health Check")

    try:
        response = make_request("GET", "/api/v1/embeddings/health")
        print_response(response)

        if response.status_code == 200:
            print("\n✅ Embeddings health check passed!")
        else:
            print("\n❌ Embeddings health check failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")


def test_embeddings_generate():
    """Test the embeddings generation endpoint."""
    print_section("TEST: Generate Embeddings")

    print("Test Data:")
    print(json.dumps(EMBEDDING_TEST_DATA, indent=2))
    print()

    try:
        response = make_request(
            "POST", "/api/v1/embeddings/generate", data=EMBEDDING_TEST_DATA
        )
        print_response(response)

        if response.status_code == 200:
            data = response.json()
            embeddings = data.get("embeddings", [])
            print(f"\n✅ Embeddings generation passed!")
            print(f"   - Model: {data.get('model')}")
            print(f"   - Dimensions: {data.get('dimensions')}")
            print(f"   - Number of embeddings: {len(embeddings)}")

            # Print shape of embeddings
            if embeddings:
                print(f"\n   Embedding shapes:")
                for i, emb in enumerate(embeddings):
                    print(f"     Text {i + 1}: {len(emb)} dimensions")
                    print(f"       First 5 values: {emb[:5]}")
        else:
            print("\n❌ Embeddings generation failed!")
    except Exception as e:
        print(f"\n❌ Test failed with error: {str(e)}")


def test_error_handling():
    """Test error handling with invalid data."""
    print_section("TEST: Error Handling")

    print("Testing with invalid subtitle URL...")
    invalid_subtitle_data = {"url": "https://invalid-url.com/video", "lang": "en"}

    try:
        response = make_request(
            "POST", "/api/v1/subtitles/extract", data=invalid_subtitle_data
        )
        print_response(response)

        if response.status_code >= 400:
            print("\n✅ Error handling works correctly!")
        else:
            print("\n⚠️  Expected error response but got success!")
    except Exception as e:
        print(f"\n✅ Error handling works: {str(e)}")

    print("\n" + "-" * 80 + "\n")
    print("Testing with empty embeddings text list...")
    invalid_embedding_data = {"texts": [], "normalize": True}

    try:
        response = make_request(
            "POST", "/api/v1/embeddings/generate", data=invalid_embedding_data
        )
        print_response(response)

        if response.status_code >= 400:
            print("\n✅ Error handling works correctly!")
        else:
            print("\n⚠️  Expected error response but got success!")
    except Exception as e:
        print(f"\n✅ Error handling works: {str(e)}")


# ============================================================================
# MAIN TEST RUNNER
# ============================================================================


def run_all_tests():
    """Run all test functions."""
    print("\n")
    print("+" + "=" * 78 + "+")
    print("|" + " " * 20 + "MULTI-UTILITY SERVER API TESTS" + " " * 28 + "|")
    print("+" + "=" * 78 + "+")
    print(f"\nBase URL: {BASE_URL}")
    print(
        f"API Key: {'*' * len(API_KEY) if API_KEY != 'your-api-key-here' else 'NOT SET'}"
    )

    if API_KEY == "your-api-key-here":
        print("\n⚠️  WARNING: Please set your API key in the API_KEY variable!")
        print("   Some tests may fail without a valid API key.\n")

    # Run all tests
    test_health_check()
    test_root_endpoint()
    test_subtitles_health()
    test_subtitles_extract()
    test_audio_transcription()
    test_embeddings_health()
    test_embeddings_generate()
    test_error_handling()

    # Summary
    print_section("TEST SUMMARY")
    print("All tests completed!")
    print(
        "\nTo modify test inputs, edit the TEST_DATA section at the top of this file."
    )
    print("\nAvailable endpoints:")
    print(f"  - GET  {BASE_URL}/health")
    print(f"  - GET  {BASE_URL}/")
    print(f"  - GET  {BASE_URL}/api/v1/subtitles/health")
    print(f"  - POST {BASE_URL}/api/v1/subtitles/extract (requires network access)")
    print(f"  - POST {BASE_URL}/api/v1/subtitles/transcribe (upload audio file)")
    print(f"  - GET  {BASE_URL}/api/v1/embeddings/health")
    print(f"  - POST {BASE_URL}/api/v1/embeddings/generate")
    print()


if __name__ == "__main__":
    run_all_tests()