243 lines
8.9 KiB
Python
243 lines
8.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
Test script for OKX data collector.
|
||
|
|
|
||
|
|
This script tests the OKX collector implementation by running a single collector
|
||
|
|
for a specified trading pair and monitoring the data collection for a short period.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import sys
|
||
|
|
import signal
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
# Add project root to Python path
|
||
|
|
project_root = Path(__file__).parent.parent
|
||
|
|
sys.path.insert(0, str(project_root))
|
||
|
|
|
||
|
|
from data.exchanges.okx import OKXCollector
|
||
|
|
from data.collector_manager import CollectorManager
|
||
|
|
from data.base_collector import DataType
|
||
|
|
from utils.logger import get_logger
|
||
|
|
from database.connection import init_database
|
||
|
|
|
||
|
|
# Global shutdown flag
|
||
|
|
shutdown_flag = asyncio.Event()
|
||
|
|
|
||
|
|
def signal_handler(signum, frame):
|
||
|
|
"""Handle shutdown signals."""
|
||
|
|
print(f"\nReceived signal {signum}, shutting down...")
|
||
|
|
shutdown_flag.set()
|
||
|
|
|
||
|
|
async def test_single_collector():
|
||
|
|
"""Test a single OKX collector."""
|
||
|
|
logger = get_logger("test_okx_collector", verbose=True)
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Initialize database
|
||
|
|
logger.info("Initializing database connection...")
|
||
|
|
db_manager = init_database()
|
||
|
|
logger.info("Database initialized successfully")
|
||
|
|
|
||
|
|
# Create OKX collector for BTC-USDT
|
||
|
|
symbol = "BTC-USDT"
|
||
|
|
data_types = [DataType.TRADE, DataType.ORDERBOOK]
|
||
|
|
|
||
|
|
logger.info(f"Creating OKX collector for {symbol}")
|
||
|
|
collector = OKXCollector(
|
||
|
|
symbol=symbol,
|
||
|
|
data_types=data_types,
|
||
|
|
auto_restart=True,
|
||
|
|
health_check_interval=30.0,
|
||
|
|
store_raw_data=True
|
||
|
|
)
|
||
|
|
|
||
|
|
# Start the collector
|
||
|
|
logger.info("Starting OKX collector...")
|
||
|
|
success = await collector.start()
|
||
|
|
|
||
|
|
if not success:
|
||
|
|
logger.error("Failed to start OKX collector")
|
||
|
|
return False
|
||
|
|
|
||
|
|
logger.info("OKX collector started successfully")
|
||
|
|
|
||
|
|
# Monitor for a short period
|
||
|
|
test_duration = 60 # seconds
|
||
|
|
logger.info(f"Monitoring collector for {test_duration} seconds...")
|
||
|
|
|
||
|
|
start_time = asyncio.get_event_loop().time()
|
||
|
|
while not shutdown_flag.is_set():
|
||
|
|
# Check if test duration elapsed
|
||
|
|
elapsed = asyncio.get_event_loop().time() - start_time
|
||
|
|
if elapsed >= test_duration:
|
||
|
|
logger.info(f"Test duration ({test_duration}s) completed")
|
||
|
|
break
|
||
|
|
|
||
|
|
# Print status every 10 seconds
|
||
|
|
if int(elapsed) % 10 == 0 and int(elapsed) > 0:
|
||
|
|
status = collector.get_status()
|
||
|
|
logger.info(f"Collector status: {status['status']} - "
|
||
|
|
f"Messages: {status.get('messages_processed', 0)} - "
|
||
|
|
f"Errors: {status.get('errors', 0)}")
|
||
|
|
|
||
|
|
await asyncio.sleep(1)
|
||
|
|
|
||
|
|
# Stop the collector
|
||
|
|
logger.info("Stopping OKX collector...")
|
||
|
|
await collector.stop()
|
||
|
|
logger.info("OKX collector stopped")
|
||
|
|
|
||
|
|
# Print final statistics
|
||
|
|
final_status = collector.get_status()
|
||
|
|
logger.info("=== Final Statistics ===")
|
||
|
|
logger.info(f"Status: {final_status['status']}")
|
||
|
|
logger.info(f"Messages processed: {final_status.get('messages_processed', 0)}")
|
||
|
|
logger.info(f"Errors: {final_status.get('errors', 0)}")
|
||
|
|
logger.info(f"WebSocket state: {final_status.get('websocket_state', 'unknown')}")
|
||
|
|
|
||
|
|
if 'websocket_stats' in final_status:
|
||
|
|
ws_stats = final_status['websocket_stats']
|
||
|
|
logger.info(f"WebSocket messages received: {ws_stats.get('messages_received', 0)}")
|
||
|
|
logger.info(f"WebSocket messages sent: {ws_stats.get('messages_sent', 0)}")
|
||
|
|
logger.info(f"Pings sent: {ws_stats.get('pings_sent', 0)}")
|
||
|
|
logger.info(f"Pongs received: {ws_stats.get('pongs_received', 0)}")
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error in test: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def test_collector_manager():
|
||
|
|
"""Test multiple collectors using CollectorManager."""
|
||
|
|
logger = get_logger("test_collector_manager", verbose=True)
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Initialize database
|
||
|
|
logger.info("Initializing database connection...")
|
||
|
|
db_manager = init_database()
|
||
|
|
logger.info("Database initialized successfully")
|
||
|
|
|
||
|
|
# Create collector manager
|
||
|
|
manager = CollectorManager(
|
||
|
|
manager_name="test_manager",
|
||
|
|
global_health_check_interval=30.0
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create multiple collectors
|
||
|
|
symbols = ["BTC-USDT", "ETH-USDT", "SOL-USDT"]
|
||
|
|
collectors = []
|
||
|
|
|
||
|
|
for symbol in symbols:
|
||
|
|
logger.info(f"Creating collector for {symbol}")
|
||
|
|
collector = OKXCollector(
|
||
|
|
symbol=symbol,
|
||
|
|
data_types=[DataType.TRADE, DataType.ORDERBOOK],
|
||
|
|
auto_restart=True,
|
||
|
|
health_check_interval=30.0,
|
||
|
|
store_raw_data=True
|
||
|
|
)
|
||
|
|
collectors.append(collector)
|
||
|
|
manager.add_collector(collector)
|
||
|
|
|
||
|
|
# Start the manager
|
||
|
|
logger.info("Starting collector manager...")
|
||
|
|
success = await manager.start()
|
||
|
|
|
||
|
|
if not success:
|
||
|
|
logger.error("Failed to start collector manager")
|
||
|
|
return False
|
||
|
|
|
||
|
|
logger.info("Collector manager started successfully")
|
||
|
|
|
||
|
|
# Monitor for a short period
|
||
|
|
test_duration = 90 # seconds
|
||
|
|
logger.info(f"Monitoring collectors for {test_duration} seconds...")
|
||
|
|
|
||
|
|
start_time = asyncio.get_event_loop().time()
|
||
|
|
while not shutdown_flag.is_set():
|
||
|
|
# Check if test duration elapsed
|
||
|
|
elapsed = asyncio.get_event_loop().time() - start_time
|
||
|
|
if elapsed >= test_duration:
|
||
|
|
logger.info(f"Test duration ({test_duration}s) completed")
|
||
|
|
break
|
||
|
|
|
||
|
|
# Print status every 15 seconds
|
||
|
|
if int(elapsed) % 15 == 0 and int(elapsed) > 0:
|
||
|
|
status = manager.get_status()
|
||
|
|
stats = status.get('statistics', {})
|
||
|
|
logger.info(f"Manager status: Running={stats.get('running_collectors', 0)}, "
|
||
|
|
f"Failed={stats.get('failed_collectors', 0)}, "
|
||
|
|
f"Total={status['total_collectors']}")
|
||
|
|
|
||
|
|
# Print individual collector status
|
||
|
|
for collector_name in manager.list_collectors():
|
||
|
|
collector_status = manager.get_collector_status(collector_name)
|
||
|
|
if collector_status:
|
||
|
|
collector_info = collector_status.get('status', {})
|
||
|
|
logger.info(f" {collector_name}: {collector_info.get('status', 'unknown')} - "
|
||
|
|
f"Messages: {collector_info.get('messages_processed', 0)}")
|
||
|
|
|
||
|
|
await asyncio.sleep(1)
|
||
|
|
|
||
|
|
# Stop the manager
|
||
|
|
logger.info("Stopping collector manager...")
|
||
|
|
await manager.stop()
|
||
|
|
logger.info("Collector manager stopped")
|
||
|
|
|
||
|
|
# Print final statistics
|
||
|
|
final_status = manager.get_status()
|
||
|
|
stats = final_status.get('statistics', {})
|
||
|
|
logger.info("=== Final Manager Statistics ===")
|
||
|
|
logger.info(f"Total collectors: {final_status['total_collectors']}")
|
||
|
|
logger.info(f"Running collectors: {stats.get('running_collectors', 0)}")
|
||
|
|
logger.info(f"Failed collectors: {stats.get('failed_collectors', 0)}")
|
||
|
|
logger.info(f"Restarts performed: {stats.get('restarts_performed', 0)}")
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error in collector manager test: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def main():
|
||
|
|
"""Main test function."""
|
||
|
|
# Setup signal handlers
|
||
|
|
signal.signal(signal.SIGINT, signal_handler)
|
||
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
||
|
|
|
||
|
|
logger = get_logger("main", verbose=True)
|
||
|
|
logger.info("Starting OKX collector tests...")
|
||
|
|
|
||
|
|
# Choose test mode
|
||
|
|
test_mode = sys.argv[1] if len(sys.argv) > 1 else "single"
|
||
|
|
|
||
|
|
if test_mode == "single":
|
||
|
|
logger.info("Running single collector test...")
|
||
|
|
success = await test_single_collector()
|
||
|
|
elif test_mode == "manager":
|
||
|
|
logger.info("Running collector manager test...")
|
||
|
|
success = await test_collector_manager()
|
||
|
|
else:
|
||
|
|
logger.error(f"Unknown test mode: {test_mode}")
|
||
|
|
logger.info("Usage: python test_okx_collector.py [single|manager]")
|
||
|
|
return False
|
||
|
|
|
||
|
|
if success:
|
||
|
|
logger.info("Test completed successfully!")
|
||
|
|
else:
|
||
|
|
logger.error("Test failed!")
|
||
|
|
|
||
|
|
return success
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
try:
|
||
|
|
success = asyncio.run(main())
|
||
|
|
sys.exit(0 if success else 1)
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
print("\nTest interrupted by user")
|
||
|
|
sys.exit(1)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Test failed with error: {e}")
|
||
|
|
sys.exit(1)
|