test_cancellation.py

#!/usr/bin/env python3
"""
Test script for HTTP request cancellation functionality.
This script demonstrates how client disconnection cancels ongoing requests.
"""

import asyncio
import httpx
import json
import time

async def test_non_streaming_cancellation():
    """Test cancellation for non-streaming requests."""
    print("🧪 Testing non-streaming request cancellation...")
    
    async with httpx.AsyncClient(timeout=30) as client:
        try:
            # Start a long-running request
            task = asyncio.create_task(
                client.post(
                    "http://localhost:8082/v1/messages",
                    json={
                        "model": "claude-3-5-sonnet-20241022",
                        "max_tokens": 1000,
                        "messages": [
                            {"role": "user", "content": "Write a very long story about a journey through space that takes at least 500 words."}
                        ]
                    }
                )
            )
            
            # Cancel after 2 seconds
            await asyncio.sleep(2)
            task.cancel()
            
            try:
                await task
                print("❌ Request should have been cancelled")
            except asyncio.CancelledError:
                print("✅ Non-streaming request cancelled successfully")
                
        except Exception as e:
            print(f"❌ Non-streaming test error: {e}")

async def test_streaming_cancellation():
    """Test cancellation for streaming requests."""
    print("\n🧪 Testing streaming request cancellation...")
    
    async with httpx.AsyncClient(timeout=30) as client:
        try:
            # Start streaming request
            async with client.stream(
                "POST",
                "http://localhost:8082/v1/messages",
                json={
                    "model": "claude-3-5-sonnet-20241022",
                    "max_tokens": 1000,
                    "messages": [
                        {"role": "user", "content": "Write a very long story about a journey through space that takes at least 500 words."}
                    ],
                    "stream": True
                }
            ) as response:
                if response.status_code == 200:
                    print("✅ Streaming request started successfully")
                    
                    # Read a few chunks then simulate client disconnect
                    chunk_count = 0
                    async for line in response.aiter_lines():
                        if line.strip():
                            chunk_count += 1
                            print(f"📦 Received chunk {chunk_count}: {line[:100]}...")
                            
                            # Simulate client disconnect after 3 chunks
                            if chunk_count >= 3:
                                print("🔌 Simulating client disconnect...")
                                break
                    
                    print("✅ Streaming request cancelled successfully")
                else:
                    print(f"❌ Streaming request failed: {response.status_code}")
                    
        except Exception as e:
            print(f"❌ Streaming test error: {e}")

async def test_server_running():
    """Test if the server is running."""
    print("🔍 Checking if server is running...")
    
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            response = await client.get("http://localhost:8082/health")
            if response.status_code == 200:
                print("✅ Server is running and healthy")
                return True
            else:
                print(f"❌ Server health check failed: {response.status_code}")
                return False
    except Exception as e:
        print(f"❌ Cannot connect to server: {e}")
        print("💡 Make sure to start the server with: python start_proxy.py")
        return False

async def main():
    """Main test function."""
    print("🚀 Starting HTTP request cancellation tests")
    print("=" * 50)
    
    # Check if server is running
    if not await test_server_running():
        return
    
    print("\n" + "=" * 50)
    
    # Test non-streaming cancellation
    await test_non_streaming_cancellation()
    
    # Test streaming cancellation  
    await test_streaming_cancellation()
    
    print("\n" + "=" * 50)
    print("✅ All cancellation tests completed!")
    print("\n💡 Note: The actual cancellation behavior depends on:")
    print("   - Client implementation (httpx in this case)")
    print("   - Network conditions")
    print("   - Server response to client disconnection")
    print("   - Whether the underlying OpenAI API supports cancellation")

if __name__ == "__main__":
    asyncio.run(main())