validate-system.py

#!/usr/bin/env python3

import requests
import json
import time
import sys
from typing import Dict, Any, List
from dataclasses import dataclass

@dataclass
class TestResult:
    name: str
    passed: bool
    message: str
    response: Any = None

class SystemValidator:
    def __init__(self):
        self.base_url = "http://localhost"
        self.router_url = f"{self.base_url}:8000"
        self.math_url = f"{self.base_url}:8001"
        self.weather_url = f"{self.base_url}:8002"
        self.qna_url = f"{self.base_url}:8003"
        
        self.session = requests.Session()
        self.session.timeout = 10
    
    def wait_for_services(self, max_wait=30):
        print("⏳ Waiting for services to be ready...")
        
        services = [
            ("Router", self.router_url),
            ("Math Agent", self.math_url),
            ("Weather Agent", self.weather_url),
            ("QnA Agent", self.qna_url)
        ]
        
        for name, url in services:
            ready = False
            for _ in range(max_wait):
                try:
                    response = self.session.get(f"{url}/health")
                    if response.status_code == 200:
                        ready = True
                        break
                except:
                    pass
                time.sleep(1)
            
            if not ready:
                print(f"❌ {name} service not ready at {url}")
                return False
        
        return True
    
    def run_all_tests(self) -> List[TestResult]:
        results = []
        
        results.extend(self.test_health_endpoints())
        results.extend(self.test_direct_agents())
        results.extend(self.test_router_routing())
        results.extend(self.test_edge_cases())
        
        return results
    
    def test_health_endpoints(self) -> List[TestResult]:
        results = []
        
        print("\n🏥 Testing Health Endpoints...")
        
        services = [
            ("Router", self.router_url),
            ("Math Agent", self.math_url),
            ("Weather Agent", self.weather_url),
            ("QnA Agent", self.qna_url)
        ]
        
        for name, url in services:
            try:
                response = self.session.get(f"{url}/health")
                if response.status_code == 200:
                    data = response.json()
                    results.append(TestResult(
                        name=f"{name} Health Check",
                        passed=True,
                        message=f"Service healthy: {data.get('status', 'unknown')}",
                        response=data
                    ))
                else:
                    results.append(TestResult(
                        name=f"{name} Health Check",
                        passed=False,
                        message=f"HTTP {response.status_code}: {response.text}"
                    ))
            except Exception as e:
                results.append(TestResult(
                    name=f"{name} Health Check",
                    passed=False,
                    message=f"Connection failed: {str(e)}"
                ))
        
        return results
    
    def test_direct_agents(self) -> List[TestResult]:
        results = []
        
        print("\n🎯 Testing Direct Agent APIs...")
        
        # Test Math Agent
        try:
            response = self.session.post(f"{self.math_url}/solve", json={"question": "2 + 2"})
            if response.status_code == 200:
                data = response.json()
                results.append(TestResult(
                    name="Math Agent Direct API",
                    passed=data.get("result") == 4.0,
                    message=f"Math calculation: {data.get('expression')} = {data.get('result')}",
                    response=data
                ))
            else:
                results.append(TestResult(
                    name="Math Agent Direct API",
                    passed=False,
                    message=f"HTTP {response.status_code}: {response.text}"
                ))
        except Exception as e:
            results.append(TestResult(
                name="Math Agent Direct API",
                passed=False,
                message=f"Error: {str(e)}"
            ))
        
        # Test Weather Agent
        try:
            response = self.session.post(f"{self.weather_url}/get", json={"question": "weather in Paris"})
            if response.status_code == 200:
                data = response.json()
                results.append(TestResult(
                    name="Weather Agent Direct API",
                    passed="weather" in data and "city" in data,
                    message=f"Weather for {data.get('city')}: {data.get('weather', {}).get('condition')}",
                    response=data
                ))
            else:
                results.append(TestResult(
                    name="Weather Agent Direct API",
                    passed=False,
                    message=f"HTTP {response.status_code}: {response.text}"
                ))
        except Exception as e:
            results.append(TestResult(
                name="Weather Agent Direct API",
                passed=False,
                message=f"Error: {str(e)}"
            ))
        
        # Test QnA Agent
        try:
            response = self.session.post(f"{self.qna_url}/ask", json={"question": "Hello"})
            if response.status_code == 200:
                data = response.json()
                results.append(TestResult(
                    name="QnA Agent Direct API",
                    passed="response" in data,
                    message=f"QnA response: {data.get('response', 'No response')}",
                    response=data
                ))
            else:
                results.append(TestResult(
                    name="QnA Agent Direct API",
                    passed=False,
                    message=f"HTTP {response.status_code}: {response.text}"
                ))
        except Exception as e:
            results.append(TestResult(
                name="QnA Agent Direct API",
                passed=False,
                message=f"Error: {str(e)}"
            ))
        
        return results
    
    def test_router_routing(self) -> List[TestResult]:
        results = []
        
        print("\n🔀 Testing Router Intelligence...")
        
        test_cases = [
            {
                "question": "What is 15 * 8 + 7?",
                "expected_agent": "math",
                "description": "Complex math expression"
            },
            {
                "question": "Calculate the square root of 144",
                "expected_agent": "math",
                "description": "Math function"
            },
            {
                "question": "What's the weather like in Tokyo?",
                "expected_agent": "weather",
                "description": "Weather query with city"
            },
            {
                "question": "Is it raining in London?",
                "expected_agent": "weather",
                "description": "Weather condition query"
            },
            {
                "question": "How hot is it in Dubai?",
                "expected_agent": "weather",
                "description": "Temperature query"
            },
            {
                "question": "What is the capital of France?",
                "expected_agent": "qna",
                "description": "General knowledge question"
            },
            {
                "question": "Hello, how are you?",
                "expected_agent": "qna",
                "description": "Conversational query"
            },
            {
                "question": "Tell me about artificial intelligence",
                "expected_agent": "qna",
                "description": "General information request"
            }
        ]
        
        for test_case in test_cases:
            try:
                response = self.session.post(
                    f"{self.router_url}/ask",
                    json={"question": test_case["question"]}
                )
                
                if response.status_code == 200:
                    data = response.json()
                    routed_agent = data.get("agent", "").lower()
                    expected = test_case["expected_agent"]
                    
                    passed = expected in routed_agent
                    results.append(TestResult(
                        name=f"Routing: {test_case['description']}",
                        passed=passed,
                        message=f"Expected: {expected}, Got: {routed_agent}",
                        response=data
                    ))
                else:
                    results.append(TestResult(
                        name=f"Routing: {test_case['description']}",
                        passed=False,
                        message=f"HTTP {response.status_code}: {response.text}"
                    ))
            except Exception as e:
                results.append(TestResult(
                    name=f"Routing: {test_case['description']}",
                    passed=False,
                    message=f"Error: {str(e)}"
                ))
        
        return results
    
    def test_edge_cases(self) -> List[TestResult]:
        results = []
        
        print("\n⚠️  Testing Edge Cases...")
        
        edge_cases = [
            {
                "question": "",
                "description": "Empty question",
                "expect_error": True
            },
            {
                "question": "   ",
                "description": "Whitespace only question",
                "expect_error": True
            },
            {
                "question": "a" * 1000,
                "description": "Very long question",
                "expect_error": False
            }
        ]
        
        for case in edge_cases:
            try:
                response = self.session.post(
                    f"{self.router_url}/ask",
                    json={"question": case["question"]}
                )
                
                if case["expect_error"]:
                    passed = response.status_code >= 400
                    message = f"Expected error, got HTTP {response.status_code}"
                else:
                    passed = response.status_code == 200
                    message = f"Got HTTP {response.status_code}"
                
                results.append(TestResult(
                    name=f"Edge Case: {case['description']}",
                    passed=passed,
                    message=message,
                    response=response.json() if response.status_code == 200 else response.text
                ))
            except Exception as e:
                results.append(TestResult(
                    name=f"Edge Case: {case['description']}",
                    passed=False,
                    message=f"Error: {str(e)}"
                ))
        
        return results
    
    def print_results(self, results: List[TestResult]):
        print("\n" + "="*50)
        print("📊 VALIDATION RESULTS")
        print("="*50)
        
        passed_count = 0
        total_count = len(results)
        
        for result in results:
            status = "✅" if result.passed else "❌"
            print(f"{status} {result.name}: {result.message}")
            if result.passed:
                passed_count += 1
        
        print("\n" + "-"*50)
        print(f"🎯 SUMMARY: {passed_count}/{total_count} tests passed ({passed_count/total_count*100:.1f}%)")
        
        if passed_count == total_count:
            print("🎉 All tests passed! System is fully functional.")
            return True
        else:
            print("⚠️  Some tests failed. Please check the issues above.")
            return False

def main():
    validator = SystemValidator()
    
    print("🧪 AI AnswerBot System Validation")
    print("="*50)
    
    if not validator.wait_for_services():
        print("❌ Services are not ready. Please ensure all containers are running.")
        sys.exit(1)
    
    results = validator.run_all_tests()
    
    success = validator.print_results(results)
    
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()