test_deployment.py

#!/usr/bin/env python3
"""
Python test script for Hugging Face Space deployment
Usage: python test_deployment.py [space-url]
"""

import sys
import requests
import json
from typing import Dict, Any

# ANSI color codes
GREEN = '\033[0;32m'
RED = '\033[0;31m'
YELLOW = '\033[1;33m'
BLUE = '\033[0;34m'
NC = '\033[0m'  # No Color


def print_test(test_name: str):
    """Print test header"""
    print(f"\n{YELLOW}{test_name}{NC}")


def print_result(passed: bool, message: str, response: Any = None):
    """Print test result"""
    status = f"{GREEN}✓ PASSED{NC}" if passed else f"{RED}✗ FAILED{NC}"
    print(f"{status} - {message}")
    if response:
        try:
            print(f"Response: {json.dumps(response, indent=2)}")
        except:
            print(f"Response: {response}")


def test_health(base_url: str) -> bool:
    """Test health endpoint"""
    print_test("Test 1: Health Check")
    print(f"GET {base_url}/health")
    
    try:
        response = requests.get(f"{base_url}/health", timeout=10)
        passed = response.status_code == 200
        print_result(
            passed,
            f"Status: {response.status_code}",
            response.json() if passed else response.text
        )
        return passed
    except Exception as e:
        print_result(False, f"Error: {str(e)}")
        return False


def test_service_info(base_url: str) -> bool:
    """Test root endpoint"""
    print_test("Test 2: Service Info")
    print(f"GET {base_url}/")
    
    try:
        response = requests.get(f"{base_url}/", timeout=10)
        passed = response.status_code == 200
        print_result(
            passed,
            f"Status: {response.status_code}",
            response.json() if passed else response.text
        )
        return passed
    except Exception as e:
        print_result(False, f"Error: {str(e)}")
        return False


def test_list_tools(base_url: str) -> bool:
    """Test tools listing endpoint"""
    print_test("Test 3: List Available Tools")
    print(f"GET {base_url}/tools")
    
    try:
        response = requests.get(f"{base_url}/tools", timeout=10)
        passed = response.status_code == 200
        print_result(
            passed,
            f"Status: {response.status_code}",
            response.json() if passed else response.text
        )
        return passed
    except Exception as e:
        print_result(False, f"Error: {str(e)}")
        return False


def test_youtube_search(base_url: str) -> bool:
    """Test YouTube search endpoint"""
    print_test("Test 4: YouTube Video Search")
    print(f"POST {base_url}/search")
    
    payload = {
        "query": "Python programming",
        "max_results": 2
    }
    
    try:
        response = requests.post(
            f"{base_url}/search",
            json=payload,
            timeout=15
        )
        
        if response.status_code == 200:
            print_result(True, f"Status: {response.status_code}", response.json())
            return True
        elif response.status_code == 500 and "YOUTUBE_API_KEY" in response.text:
            print(f"{YELLOW}⚠ API KEY NOT SET{NC} - Status: {response.status_code}")
            print("Please set YOUTUBE_API_KEY in your Space settings")
            print(f"Response: {response.json()}")
            return False
        else:
            print_result(False, f"Status: {response.status_code}", response.text)
            return False
    except Exception as e:
        print_result(False, f"Error: {str(e)}")
        return False


def test_error_handling(base_url: str) -> bool:
    """Test error handling with missing query"""
    print_test("Test 5: Error Handling (Missing Query)")
    print(f"POST {base_url}/search (with empty query)")
    
    payload = {}
    
    try:
        response = requests.post(
            f"{base_url}/search",
            json=payload,
            timeout=10
        )
        
        # We expect a 400 error
        passed = response.status_code == 400
        print_result(
            passed,
            f"Status: {response.status_code} (Expected 400 error)",
            response.json() if response.status_code == 400 else response.text
        )
        return passed
    except Exception as e:
        print_result(False, f"Error: {str(e)}")
        return False


def main():
    """Run all tests"""
    # Get Space URL from command line or use default
    base_url = sys.argv[1] if len(sys.argv) > 1 else "https://ocx2025-basicsearch.hf.space"
    
    # Remove trailing slash if present
    base_url = base_url.rstrip('/')
    
    print("=" * 50)
    print(f"Testing Hugging Face Space: {base_url}")
    print("=" * 50)
    
    # Run all tests
    results = []
    results.append(("Health Check", test_health(base_url)))
    results.append(("Service Info", test_service_info(base_url)))
    results.append(("List Tools", test_list_tools(base_url)))
    results.append(("YouTube Search", test_youtube_search(base_url)))
    results.append(("Error Handling", test_error_handling(base_url)))
    
    # Print summary
    print("\n" + "=" * 50)
    print("Test Summary")
    print("=" * 50)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = f"{GREEN}✓{NC}" if result else f"{RED}✗{NC}"
        print(f"{status} {test_name}")
    
    print(f"\n{BLUE}Results: {passed}/{total} tests passed{NC}")
    print("=" * 50)
    
    # Exit with appropriate code
    sys.exit(0 if passed == total else 1)


if __name__ == "__main__":
    main()