test_api.py

"""
Comprehensive API Testing Script
Tests all endpoints with realistic medical scenarios
"""

import requests
import json
import time
from typing import Dict, List
import os


class HealthAPITester:
    """Test suite for AI Health Diagnostics API"""
    
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url
        self.user_ids = []
        self.test_results = []
    
    def print_section(self, title: str):
        """Print formatted section header"""
        print("\n" + "=" * 70)
        print(f"  {title}")
        print("=" * 70)
    
    def print_result(self, test_name: str, passed: bool, details: str = ""):
        """Print test result"""
        status = "✅ PASS" if passed else "❌ FAIL"
        print(f"{status} | {test_name}")
        if details:
            print(f"         {details}")
        self.test_results.append((test_name, passed))
    
    def test_health_check(self):
        """Test root endpoint"""
        self.print_section("1. Health Check")
        
        try:
            response = requests.get(f"{self.base_url}/")
            passed = response.status_code == 200
            details = response.json().get("status", "No status")
            self.print_result("API Health Check", passed, details)
            return passed
        except Exception as e:
            self.print_result("API Health Check", False, str(e))
            return False
    
    def test_create_profiles(self):
        """Test profile creation with various user types"""
        self.print_section("2. User Profile Creation")
        
        test_profiles = [
            {
                "name": "Young Adult - No Conditions",
                "data": {
                    "age": 25,
                    "gender": "female",
                    "allergies": "",
                    "conditions": ""
                }
            },
            {
                "name": "Middle-aged - Multiple Allergies",
                "data": {
                    "age": 45,
                    "gender": "male",
                    "allergies": "penicillin, shellfish",
                    "conditions": "hypertension"
                }
            },
            {
                "name": "Senior - Chronic Conditions",
                "data": {
                    "age": 68,
                    "gender": "female",
                    "allergies": "sulfa drugs",
                    "conditions": "diabetes, arthritis"
                }
            },
            {
                "name": "Child Profile",
                "data": {
                    "age": 8,
                    "gender": "male",
                    "allergies": "peanuts",
                    "conditions": "asthma"
                }
            }
        ]
        
        for profile in test_profiles:
            try:
                response = requests.post(
                    f"{self.base_url}/create-profile",
                    json=profile["data"]
                )
                
                if response.status_code == 200:
                    result = response.json()
                    user_id = result["data"]["user_id"]
                    self.user_ids.append({
                        "name": profile["name"],
                        "id": user_id,
                        "profile": profile["data"]
                    })
                    self.print_result(
                        f"Create Profile: {profile['name']}",
                        True,
                        f"User ID: {user_id[:8]}..."
                    )
                else:
                    self.print_result(
                        f"Create Profile: {profile['name']}",
                        False,
                        f"Status: {response.status_code}"
                    )
            except Exception as e:
                self.print_result(
                    f"Create Profile: {profile['name']}",
                    False,
                    str(e)
                )
    
    def test_symptom_analysis(self):
        """Test symptom checking with realistic scenarios"""
        self.print_section("3. Symptom Analysis")
        
        test_cases = [
            {
                "name": "Common Cold Symptoms",
                "symptoms": "runny nose, cough, sore throat, mild fever, fatigue",
                "expected_conditions": ["Common Cold", "Influenza"]
            },
            {
                "name": "Flu-like Symptoms",
                "symptoms": "high fever, severe headache, body aches, cough, fatigue",
                "expected_conditions": ["Influenza", "COVID-19"]
            },
            {
                "name": "Food Poisoning",
                "symptoms": "nausea, vomiting, diarrhea, abdominal cramps",
                "expected_conditions": ["Gastroenteritis", "Food Poisoning"]
            },
            {
                "name": "Migraine Symptoms",
                "symptoms": "severe headache, sensitivity to light, nausea",
                "expected_conditions": ["Migraine"]
            },
            {
                "name": "Allergic Reaction",
                "symptoms": "rash, itching, swelling, redness",
                "expected_conditions": ["Allergic Reaction"]
            },
            {
                "name": "Respiratory Infection",
                "symptoms": "persistent cough, shortness of breath, chest pain, fever",
                "expected_conditions": ["Pneumonia", "Bronchitis"]
            },
            {
                "name": "Emergency Symptoms",
                "symptoms": "severe chest pain, difficulty breathing, confusion",
                "expected_conditions": ["Cardiac Issue", "Heart Condition"]
            }
        ]
        
        if not self.user_ids:
            self.print_result("Symptom Analysis", False, "No user profiles available")
            return
        
        user = self.user_ids[0]  # Use first user
        
        for case in test_cases:
            try:
                response = requests.post(
                    f"{self.base_url}/symptom-check",
                    json={
                        "user_id": user["id"],
                        "symptoms": case["symptoms"]
                    }
                )
                
                if response.status_code == 200:
                    result = response.json()
                    conditions = result["data"]["possible_conditions"]
                    
                    # Check if any expected condition is in top 3 results
                    detected = [c["condition"] for c in conditions[:3]]
                    found_match = any(
                        exp in detected 
                        for exp in case["expected_conditions"]
                    )
                    
                    if conditions:
                        top_condition = conditions[0]
                        details = (
                            f"{top_condition['condition']} "
                            f"(confidence: {top_condition['confidence']:.2f}, "
                            f"urgency: {top_condition['urgency']})"
                        )
                        self.print_result(
                            f"Analyze: {case['name']}",
                            True,
                            details
                        )
                    else:
                        self.print_result(
                            f"Analyze: {case['name']}",
                            False,
                            "No conditions detected"
                        )
                else:
                    self.print_result(
                        f"Analyze: {case['name']}",
                        False,
                        f"Status: {response.status_code}"
                    )
                
                time.sleep(0.5)  # Rate limiting
                
            except Exception as e:
                self.print_result(
                    f"Analyze: {case['name']}",
                    False,
                    str(e)
                )
    
    def test_image_analysis(self):
        """Test image analysis (with synthetic images)"""
        self.print_section("4. Medical Image Analysis")
        
        # Create sample test images if they don't exist
        test_images = self._create_test_images()
        
        if not self.user_ids:
            self.print_result("Image Analysis", False, "No user profiles available")
            return
        
        user = self.user_ids[0]
        
        for image_file, image_type in test_images:
            if not os.path.exists(image_file):
                continue
            
            try:
                with open(image_file, 'rb') as f:
                    files = {'file': (image_file, f, 'image/jpeg')}
                    data = {
                        'user_id': user["id"],
                        'image_type': image_type
                    }
                    
                    response = requests.post(
                        f"{self.base_url}/analyze-image",
                        files=files,
                        data=data
                    )
                
                if response.status_code == 200:
                    result = response.json()
                    analysis = result["data"]
                    details = (
                        f"{analysis['detected_condition']} "
                        f"(confidence: {analysis['confidence']:.2f}, "
                        f"urgency: {analysis['urgency']})"
                    )
                    self.print_result(
                        f"Analyze Image: {os.path.basename(image_file)}",
                        True,
                        details
                    )
                else:
                    self.print_result(
                        f"Analyze Image: {os.path.basename(image_file)}",
                        False,
                        f"Status: {response.status_code}"
                    )
                
                time.sleep(0.5)
                
            except Exception as e:
                self.print_result(
                    f"Analyze Image: {os.path.basename(image_file)}",
                    False,
                    str(e)
                )
    
    def test_user_history(self):
        """Test retrieving user history"""
        self.print_section("5. User History Retrieval")
        
        if not self.user_ids:
            self.print_result("User History", False, "No user profiles available")
            return
        
        for user in self.user_ids[:2]:  # Test first 2 users
            try:
                response = requests.get(
                    f"{self.base_url}/user-history/{user['id']}"
                )
                
                if response.status_code == 200:
                    result = response.json()
                    data = result["data"]
                    symptom_count = len(data.get("symptom_checks", []))
                    image_count = len(data.get("image_analyses", []))
                    details = (
                        f"{symptom_count} symptom checks, "
                        f"{image_count} image analyses"
                    )
                    self.print_result(
                        f"History: {user['name']}",
                        True,
                        details
                    )
                else:
                    self.print_result(
                        f"History: {user['name']}",
                        False,
                        f"Status: {response.status_code}"
                    )
            except Exception as e:
                self.print_result(
                    f"History: {user['name']}",
                    False,
                    str(e)
                )
    
    def test_edge_cases(self):
        """Test edge cases and error handling"""
        self.print_section("6. Edge Cases & Error Handling")
        
        # Test invalid user ID
        try:
            response = requests.post(
                f"{self.base_url}/symptom-check",
                json={
                    "user_id": "invalid-user-id-12345",
                    "symptoms": "headache"
                }
            )
            passed = response.status_code == 404
            self.print_result(
                "Invalid User ID",
                passed,
                "Correctly rejected" if passed else "Should return 404"
            )
        except Exception as e:
            self.print_result("Invalid User ID", False, str(e))
        
        # Test empty symptoms
        if self.user_ids:
            try:
                response = requests.post(
                    f"{self.base_url}/symptom-check",
                    json={
                        "user_id": self.user_ids[0]["id"],
                        "symptoms": ""
                    }
                )
                result = response.json()
                passed = response.status_code == 200
                self.print_result(
                    "Empty Symptoms",
                    passed,
                    "Handled gracefully" if passed else "Should handle empty input"
                )
            except Exception as e:
                self.print_result("Empty Symptoms", False, str(e))
    
    def _create_test_images(self) -> List[tuple]:
        """Create simple test images"""
        try:
            from PIL import Image, ImageDraw
            import numpy as np
            
            os.makedirs("test_images", exist_ok=True)
            
            images = []
            
            # Create a simple test image
            img = Image.new('RGB', (224, 224), color='lightpink')
            draw = ImageDraw.Draw(img)
            draw.ellipse([50, 50, 174, 174], fill='red', outline='darkred')
            img.save("test_images/skin_rash.jpg")
            images.append(("test_images/skin_rash.jpg", "skin"))
            
            return images
        except:
            return []
    
    def print_summary(self):
        """Print test summary"""
        self.print_section("Test Summary")
        
        total = len(self.test_results)
        passed = sum(1 for _, result in self.test_results if result)
        failed = total - passed
        
        print(f"\nTotal Tests: {total}")
        print(f"✅ Passed: {passed}")
        print(f"❌ Failed: {failed}")
        print(f"Success Rate: {(passed/total*100) if total > 0 else 0:.1f}%")
        
        if failed > 0:
            print("\n⚠️  Failed Tests:")
            for name, result in self.test_results:
                if not result:
                    print(f"   - {name}")
        
        print("\n" + "=" * 70)
    
    def run_all_tests(self):
        """Run complete test suite"""
        print("\n🧪 AI Health Diagnostics API - Comprehensive Test Suite")
        print(f"🌐 Testing endpoint: {self.base_url}")
        
        # Run tests in order
        if not self.test_health_check():
            print("\n❌ API is not running. Start server with:")
            print("   uvicorn main:app --reload")
            return
        
        self.test_create_profiles()
        self.test_symptom_analysis()
        self.test_image_analysis()
        self.test_user_history()
        self.test_edge_cases()
        
        # Print summary
        self.print_summary()


def main():
    """Main test runner"""
    import sys
    
    base_url = "http://localhost:8000"
    if len(sys.argv) > 1:
        base_url = sys.argv[1]
    
    tester = HealthAPITester(base_url)
    tester.run_all_tests()


if __name__ == "__main__":
    main()