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