155 lines
4.9 KiB
Python
155 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple test to verify recursion fix in WebSocket task management.
|
|
"""
|
|
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add project root to path
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
from data.exchanges.okx.websocket import OKXWebSocketClient, OKXSubscription, OKXChannelType
|
|
from utils.logger import get_logger
|
|
|
|
|
|
async def test_rapid_connection_cycles():
|
|
"""Test rapid connect/disconnect cycles to verify no recursion errors."""
|
|
logger = get_logger("recursion_test", verbose=False)
|
|
|
|
print("🧪 Testing WebSocket Recursion Fix")
|
|
print("=" * 40)
|
|
|
|
for cycle in range(5):
|
|
print(f"\n🔄 Cycle {cycle + 1}/5: Rapid connect/disconnect")
|
|
|
|
ws_client = OKXWebSocketClient(
|
|
component_name=f"test_client_{cycle}",
|
|
max_reconnect_attempts=2,
|
|
logger=logger
|
|
)
|
|
|
|
try:
|
|
# Connect
|
|
success = await ws_client.connect()
|
|
if not success:
|
|
print(f" ❌ Connection failed in cycle {cycle + 1}")
|
|
continue
|
|
|
|
# Subscribe
|
|
subscriptions = [
|
|
OKXSubscription(OKXChannelType.TRADES.value, "BTC-USDT")
|
|
]
|
|
await ws_client.subscribe(subscriptions)
|
|
|
|
# Quick activity
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Disconnect (this should not cause recursion)
|
|
await ws_client.disconnect()
|
|
print(f" ✅ Cycle {cycle + 1} completed successfully")
|
|
|
|
except RecursionError as e:
|
|
print(f" ❌ Recursion error in cycle {cycle + 1}: {e}")
|
|
return False
|
|
except Exception as e:
|
|
print(f" ⚠️ Other error in cycle {cycle + 1}: {e}")
|
|
# Continue with other cycles
|
|
|
|
# Small delay between cycles
|
|
await asyncio.sleep(0.2)
|
|
|
|
print("\n✅ All cycles completed without recursion errors")
|
|
return True
|
|
|
|
|
|
async def test_concurrent_shutdowns():
|
|
"""Test concurrent client shutdowns to verify no recursion."""
|
|
logger = get_logger("concurrent_shutdown_test", verbose=False)
|
|
|
|
print("\n🔄 Testing Concurrent Shutdowns")
|
|
print("=" * 40)
|
|
|
|
# Create multiple clients
|
|
clients = []
|
|
for i in range(3):
|
|
client = OKXWebSocketClient(
|
|
component_name=f"concurrent_client_{i}",
|
|
logger=logger
|
|
)
|
|
clients.append(client)
|
|
|
|
try:
|
|
# Connect all clients
|
|
connect_tasks = [client.connect() for client in clients]
|
|
results = await asyncio.gather(*connect_tasks, return_exceptions=True)
|
|
|
|
successful_connections = sum(1 for r in results if r is True)
|
|
print(f"📡 Connected {successful_connections}/3 clients")
|
|
|
|
# Let them run briefly
|
|
await asyncio.sleep(1)
|
|
|
|
# Shutdown all concurrently (this is where recursion might occur)
|
|
print("🛑 Shutting down all clients concurrently...")
|
|
shutdown_tasks = [client.disconnect() for client in clients]
|
|
|
|
# Use wait_for to prevent hanging
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.gather(*shutdown_tasks, return_exceptions=True),
|
|
timeout=10.0
|
|
)
|
|
print("✅ All clients shut down successfully")
|
|
return True
|
|
|
|
except asyncio.TimeoutError:
|
|
print("⚠️ Shutdown timeout - but no recursion errors")
|
|
return True # Timeout is better than recursion
|
|
|
|
except RecursionError as e:
|
|
print(f"❌ Recursion error during concurrent shutdown: {e}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"⚠️ Other error during test: {e}")
|
|
return True # Other errors are acceptable for this test
|
|
|
|
|
|
async def main():
|
|
"""Run recursion fix tests."""
|
|
print("🚀 WebSocket Recursion Fix Test Suite")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
# Test 1: Rapid cycles
|
|
test1_success = await test_rapid_connection_cycles()
|
|
|
|
# Test 2: Concurrent shutdowns
|
|
test2_success = await test_concurrent_shutdowns()
|
|
|
|
# Summary
|
|
print("\n" + "=" * 50)
|
|
print("📋 Test Summary:")
|
|
print(f" Rapid Cycles: {'✅ PASS' if test1_success else '❌ FAIL'}")
|
|
print(f" Concurrent Shutdowns: {'✅ PASS' if test2_success else '❌ FAIL'}")
|
|
|
|
if test1_success and test2_success:
|
|
print("\n🎉 All tests passed! Recursion issue fixed.")
|
|
return 0
|
|
else:
|
|
print("\n❌ Some tests failed.")
|
|
return 1
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n⏹️ Tests interrupted")
|
|
return 1
|
|
except Exception as e:
|
|
print(f"\n💥 Test suite failed: {e}")
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit_code = asyncio.run(main())
|
|
sys.exit(exit_code) |