test_api.py

import requests
import json

# Test the Flask API
def test_api():
    url = "http://localhost:5001/compute"
    
    # Test data
    test_data = {
        "prompt": "What is the capital of France?",
        "max_length": 256,
        "temperature": 0.7,
        "top_p": 0.9
    }
    
    try:
        print("Testing the /compute endpoint...")
        print(f"Sending prompt: {test_data['prompt']}")
        
        response = requests.post(url, json=test_data)
        
        if response.status_code == 200:
            result = response.json()
            print("\nResponse received successfully!")
            print(f"Status: {result['status']}")
            print(f"Response: {result['response']}")
        else:
            print(f"Error: {response.status_code}")
            print(response.text)
            
    except requests.exceptions.ConnectionError:
        print("Error: Could not connect to the server. Make sure the Flask app is running on port 5001.")
    except Exception as e:
        print(f"Error: {str(e)}")

def test_health_check():
    url = "http://localhost:5001/"
    
    try:
        print("Testing health check endpoint...")
        response = requests.get(url)
        
        if response.status_code == 200:
            result = response.json()
            print("Health check successful!")
            print(json.dumps(result, indent=2))
        else:
            print(f"Error: {response.status_code}")
            print(response.text)
            
    except requests.exceptions.ConnectionError:
        print("Error: Could not connect to the server. Make sure the Flask app is running on port 5001.")
    except Exception as e:
        print(f"Error: {str(e)}")

if __name__ == "__main__":
    print("=== Flask API Test ===")
    test_health_check()
    print("\n" + "="*50 + "\n")
    test_api()