Files
claude-code-proxy/test_cancellation.py
Holegots 6ce4226111 🎉 Init Project
2025-06-23 01:27:50 +08:00

130 lines
4.8 KiB
Python

#!/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())