test_oauth_mcp.py

#!/usr/bin/env python3
"""
Test script for Multi-Cloud OAuth MCP functionality.
This demonstrates how to use the updated MCP client with authentication headers.
"""

import asyncio
from mcp.client.streamable_http import streamablehttp_client
from mcp import ClientSession

# Test configuration - replace with your actual credentials
user_token = "..."  # Replace with valid user token
user_cloud_id = "..."  # Replace with valid cloud ID

async def test_http_mcp_with_auth():
    """Test HTTP MCP connection with authentication headers."""
    print("🔧 Testing HTTP MCP with Multi-Cloud OAuth")
    print("=" * 60)
    
    # Example based on your provided sample
    async with streamablehttp_client(
        "http://localhost:9000/mcp",
        headers={
            "Authorization": f"Bearer {user_token}",
            "X-Atlassian-Cloud-Id": user_cloud_id,
        },
    ) as (read_stream, write_stream, _):
        async with ClientSession(read_stream, write_stream) as session:
            await session.initialize()
            
            print("✓ Successfully connected to HTTP MCP server")
            
            try:
                # List available tools
                print("\n📋 Available tools:")
                tools_response = await session.list_tools()
                for item in tools_response:
                    if isinstance(item, tuple) and item[0] == "tools":
                        for tool in item[1]:
                            print(f"  - {tool.name}: {tool.description}")
                
                # Test calling a tool
                print("\n🛠️  Testing tool execution:")
                response = await session.call_tool(
                    "jira_get_all_projects",
                    {},
                )
                print(f"Tool response: {response}")
                
            except Exception as e:
                print(f"Error during tool execution: {e}")


async def test_api_with_auth():
    """Test the FastAPI endpoint with authentication."""
    import httpx
    
    print("\n🌐 Testing FastAPI endpoint with authentication")
    print("=" * 60)
    
    async with httpx.AsyncClient() as client:
        try:
            # Test with authentication headers in the request body
            response = await client.post(
                "http://localhost:8000/chat",
                json={
                    "message": "List all available JIRA projects",
                    "session_id": "test-oauth-session",
                    "auth_token": user_token,
                    "cloud_id": user_cloud_id,
                    "response_format": "natural"
                }
            )
            
            if response.status_code == 200:
                result = response.json()
                print("✓ API call successful")
                print(f"Response: {result['response']}")
                print(f"Tool executed: {result['tool_executed']}")
            else:
                print(f"❌ API call failed: {response.status_code} - {response.text}")
                
        except Exception as e:
            print(f"Error calling API: {e}")


async def main():
    """Run all tests."""
    print("🔐 Multi-Cloud OAuth MCP Testing Suite")
    print("=" * 60)
    
    # Check if credentials are provided
    if user_token == "..." or user_cloud_id == "...":
        print("⚠️  Please update user_token and user_cloud_id with your actual credentials")
        print("   user_token should be your Bearer token")
        print("   user_cloud_id should be your Atlassian Cloud ID")
        return
    
    try:
        # Test 1: Direct HTTP MCP connection
        await test_http_mcp_with_auth()
        
        # Test 2: FastAPI endpoint with authentication
        await test_api_with_auth()
        
        print("\n✅ All tests completed!")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")


if __name__ == "__main__":
    print("Starting OAuth MCP tests...")
    asyncio.run(main())