test_api.py

import requests
import json

# Test the API
base_url = "http://localhost:7860"

def test_api():
    # Test root endpoint
    print("Testing root endpoint...")
    try:
        response = requests.get(f"{base_url}/")
        print(f"Root: {response.status_code} - {response.json()}")
    except Exception as e:
        print(f"Root endpoint error: {e}")
    
    # Test health endpoint
    print("\nTesting health endpoint...")
    try:
        response = requests.get(f"{base_url}/health")
        print(f"Health: {response.status_code} - {response.json()}")
    except Exception as e:
        print(f"Health endpoint error: {e}")
    
    # Test model info endpoint
    print("\nTesting model info endpoint...")
    try:
        response = requests.get(f"{base_url}/model-info")
        print(f"Model Info: {response.status_code} - {response.json()}")
    except Exception as e:
        print(f"Model info endpoint error: {e}")
    
    # Test prediction endpoint
    print("\nTesting prediction endpoint...")
    test_texts = [
        "This is a normal, safe message.",
        "I will harm you and your family!",
        "Hello, how are you doing today?",
        "This product is amazing, I love it!"
    ]
    
    for text in test_texts:
        try:
            payload = {"text": text}
            response = requests.post(f"{base_url}/predict", json=payload)
            if response.status_code == 200:
                result = response.json()
                print(f"\nText: '{text}'")
                print(f"Is Threat: {result['is_threat']}")
                print(f"Confidence: {result['final_confidence']:.3f}")
                print(f"Threat Prediction: {result['threat_prediction']:.3f}")
            else:
                print(f"Error {response.status_code}: {response.text}")
        except Exception as e:
            print(f"Prediction error for '{text}': {e}")

if __name__ == "__main__":
    test_api()