test_api_endpoints.py

#!/usr/bin/env python3
"""
Comprehensive test script for API-migrated MCP server endpoints
Tests all 8 migrated tools with sample requests
"""

import asyncio
import httpx
import json
import sysdd
from typing import Dict, Any

# Test server configuration
TEST_SERVER = "http://localhost:8000"

class APIEndpointTester:
    def __init__(self, base_url: str = TEST_SERVER):
        self.base_url = base_url
        self.client = None
        self.test_results = {}

    async def init_client(self):
        """Initialize HTTP client for testing"""
        self.client = httpx.AsyncClient(timeout=30.0)

    async def close_client(self):
        """Close HTTP client"""
        if self.client:
            await self.client.aclose()

    async def test_endpoint(self, endpoint: str, payload: Dict[Any, Any], description: str) -> bool:
        """Test a single endpoint with given payload"""
        try:
            print(f"\n🧪 Testing {endpoint}: {description}")
            
            response = await self.client.post(f"{self.base_url}{endpoint}", json=payload)
            
            if response.status_code == 200:
                result = response.json()
                
                # Check for error in response
                if "error" in result:
                    print(f"❌ {endpoint} returned error: {result['error']}")
                    return False
                
                # Basic validation
                if endpoint == "/search_funds" and "results" in result:
                    print(f"✅ {endpoint} SUCCESS: Found {result.get('count', 0)} funds")
                elif endpoint == "/postgres_query" and "response" in result:
                    print(f"✅ {endpoint} SUCCESS: Query executed")
                elif endpoint == "/describe_schema" and ("schema" in result or "tables" in result):
                    print(f"✅ {endpoint} SUCCESS: Schema described")
                elif endpoint == "/analyze_performance" and "performance_data" in result:
                    print(f"✅ {endpoint} SUCCESS: Analyzed {result.get('count', 0)} funds")
                elif endpoint == "/sector_analysis" and "summary_data" in result:
                    print(f"✅ {endpoint} SUCCESS: Sector analysis completed")
                elif endpoint == "/correlate_performance" and "correlation_matrix" in result:
                    print(f"✅ {endpoint} SUCCESS: Correlation analysis completed")
                elif endpoint == "/nav_trend_analysis" and "trend_data" in result:
                    print(f"✅ {endpoint} SUCCESS: NAV trend analysis completed")
                elif endpoint == "/risk_return_scatter" and "plot_data" in result:
                    print(f"✅ {endpoint} SUCCESS: Risk-return analysis completed")
                else:
                    print(f"⚠️ {endpoint} PARTIAL: Response format unexpected")
                    print(f"   Response keys: {list(result.keys())}")
                
                return True
            else:
                print(f"❌ {endpoint} FAILED: HTTP {response.status_code}")
                print(f"   Response: {response.text[:200]}...")
                return False
                
        except Exception as e:
            print(f"❌ {endpoint} ERROR: {str(e)}")
            return False

    async def run_all_tests(self):
        """Run comprehensive tests for all endpoints"""
        print("🚀 Starting API endpoint tests...")
        
        await self.init_client()
        
        # Test 1: Search Funds
        await self.test_endpoint(
            "/search_funds",
            {
                "fund_name": "nippon",
                "limit": 5
            },
            "Search for Nippon funds"
        )
        
        # Test 2: Postgres Query (now general API query)
        await self.test_endpoint(
            "/postgres_query",
            {
                "query": "equity funds",
                "limit": 10
            },
            "Search for equity funds"
        )
        
        # Test 3: Analyze Performance
        await self.test_endpoint(
            "/analyze_performance",
            {
                "fund_names": ["nippon", "icici"]
            },
            "Compare performance of specific funds"
        )
        
        # Test 4: Describe Schema
        await self.test_endpoint(
            "/describe_schema",
            {
                "data_type": "all"
            },
            "Get API schema information"
        )
        
        # Test 5: Sector Analysis
        await self.test_endpoint(
            "/sector_analysis",
            {
                "analysis_type": "asset_class",
                "limit": 100
            },
            "Asset class distribution analysis"
        )
        
        # Test 6: Correlate Performance
        await self.test_endpoint(
            "/correlate_performance",
            {
                "asset_class": "Equity",
                "time_period": "3year",
                "min_funds": 5
            },
            "Correlation analysis for equity funds"
        )
        
        # Test 7: NAV Trend Analysis
        await self.test_endpoint(
            "/nav_trend_analysis",
            {
                "asset_class": "Equity",
                "time_analysis": "performance_comparison",
                "limit": 10
            },
            "NAV trend analysis for equity funds"
        )
        
        # Test 8: Risk Return Scatter
        await self.test_endpoint(
            "/risk_return_scatter",
            {
                "asset_class": "Equity",
                "time_period": "3year",
                "limit": 20
            },
            "Risk-return scatter analysis"
        )
        
        # Test 9: Health Check
        await self.test_endpoint(
            "/ping",
            {},
            "Health check endpoint"
        )
        
        await self.close_client()
        
        print("\n" + "="*60)
        print("🎯 API ENDPOINT TEST SUMMARY")
        print("="*60)
        print("All core endpoints tested with sample data")
        print("✅ Migration from PostgreSQL to Scripbox API complete!")
        print("🚀 Your MCP server is ready for production use")

async def main():
    """Main test runner"""
    tester = APIEndpointTester()
    
    try:
        await tester.run_all_tests()
    except KeyboardInterrupt:
        print("\n🛑 Tests interrupted by user")
    except Exception as e:
        print(f"\n💥 Test runner error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    print("📋 MCP Server API Endpoint Tester")
    print("🔧 Make sure your server is running on http://localhost:8000")
    print("▶️ Start server with: uvicorn main:app --reload")
    print()
    
    asyncio.run(main())