test_api.py

"""
Simple test script for API endpoints.
Run this after starting the service to verify endpoints work correctly.

Usage:
    # Test local API (default port 8000)
    python test_api.py
    
    # Test local API on port 7860 (Hugging Face Spaces default)
    python test_api.py --port 7860
    
    # Test Hugging Face Spaces API
    python test_api.py --url https://bldeaw-ml-service.hf.space
"""
import requests
import json
import sys
import argparse

# Default configuration
DEFAULT_URL = "http://localhost:8000"
DEFAULT_PORT = 8000


def test_health():
    """Test health endpoint."""
    print("Testing /health endpoint...")
    response = requests.get(f"{BASE_URL}/health")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()


def test_predict():
    """Test job failure prediction endpoint."""
    print("Testing /predict/job-fail endpoint...")
    payload = {
        "zone": "prod",
        "job_nm": "daily_export_customer",
        "tasksgroup_nm": "export_group",
        "job_start_time": "2026-01-21T01:00:00",
        "duration": "01:30:00",
        "duration_sec": 5400,
        "status": "SUCCESS",
        "err_msg": "",
        "zeppelin": None,
        "explain": True
    }
    response = requests.post(
        f"{BASE_URL}/predict/job-fail",
        json=payload,
        headers={"Content-Type": "application/json"}
    )
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()


def test_anomaly():
    """Test anomaly detection endpoint."""
    print("Testing /detect/anomaly endpoint...")
    payload = {
        "features": {
            "duration_sec": 5400,
            "duration_zscore": 1.6,
            "avg_duration_7": 3000,
            "failure_rate_7": 0.15,
            "err_msg_len": 0,
            "hour_sin": 0.2588,
            "hour_cos": 0.9659
        },
        "threshold": 0.01
    }
    response = requests.post(
        f"{BASE_URL}/detect/anomaly",
        json=payload,
        headers={"Content-Type": "application/json"}
    )
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()


def main():
    """Main test function."""
    parser = argparse.ArgumentParser(description="Test API endpoints")
    parser.add_argument(
        "--url",
        type=str,
        default=None,
        help="Base URL of the API (e.g., http://localhost:8000 or https://bldeaw-ml-service.hf.space)"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=None,
        help="Port number for local API (default: 8000)"
    )
    args = parser.parse_args()
    
    # Determine base URL
    if args.url:
        base_url = args.url.rstrip('/')
    elif args.port:
        base_url = f"http://localhost:{args.port}"
    else:
        base_url = DEFAULT_URL
    
    global BASE_URL
    BASE_URL = base_url
    
    print(f"Testing API at: {BASE_URL}")
    print("=" * 60)
    print()
    
    try:
        test_health()
        test_predict()
        test_anomaly()
        print("=" * 60)
        print("✅ All tests completed successfully!")
    except requests.exceptions.ConnectionError:
        print("❌ Error: Could not connect to API.")
        print(f"   Make sure the service is running on {BASE_URL}")
        sys.exit(1)
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()