Add OKX data collector implementation and modular exchange architecture
- Introduced the `OKXCollector` and `OKXWebSocketClient` classes for real-time market data collection from the OKX exchange. - Implemented a factory pattern for creating exchange-specific collectors, enhancing modularity and scalability. - Added configuration support for the OKX collector in `config/okx_config.json`. - Updated documentation to reflect the new modular architecture and provide guidance on using the OKX collector. - Created unit tests for the OKX collector and exchange factory to ensure functionality and reliability. - Enhanced logging and error handling throughout the new implementation for improved monitoring and debugging.
This commit is contained in:
Vasily.onl
126
tests/test_exchange_factory.py
Normal file
126
tests/test_exchange_factory.py
Normal file
@@ -0,0 +1,126 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script for exchange factory pattern.
|
||||
|
||||
This script demonstrates how to use the new exchange factory
|
||||
to create collectors from different exchanges.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
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 import (
|
||||
ExchangeFactory,
|
||||
ExchangeCollectorConfig,
|
||||
create_okx_collector,
|
||||
get_supported_exchanges
|
||||
)
|
||||
from data.base_collector import DataType
|
||||
from database.connection import init_database
|
||||
from utils.logger import get_logger
|
||||
|
||||
|
||||
async def test_factory_pattern():
|
||||
"""Test the exchange factory pattern."""
|
||||
logger = get_logger("factory_test", verbose=True)
|
||||
|
||||
try:
|
||||
# Initialize database
|
||||
logger.info("Initializing database...")
|
||||
init_database()
|
||||
|
||||
# Test 1: Show supported exchanges
|
||||
logger.info("=== Supported Exchanges ===")
|
||||
supported = get_supported_exchanges()
|
||||
logger.info(f"Supported exchanges: {supported}")
|
||||
|
||||
# Test 2: Create collector using factory
|
||||
logger.info("=== Testing Exchange Factory ===")
|
||||
config = ExchangeCollectorConfig(
|
||||
exchange='okx',
|
||||
symbol='BTC-USDT',
|
||||
data_types=[DataType.TRADE, DataType.ORDERBOOK],
|
||||
auto_restart=True,
|
||||
health_check_interval=30.0,
|
||||
store_raw_data=True
|
||||
)
|
||||
|
||||
# Validate configuration
|
||||
is_valid = ExchangeFactory.validate_config(config)
|
||||
logger.info(f"Configuration valid: {is_valid}")
|
||||
|
||||
if is_valid:
|
||||
# Create collector using factory
|
||||
collector = ExchangeFactory.create_collector(config)
|
||||
logger.info(f"Created collector: {type(collector).__name__}")
|
||||
logger.info(f"Collector symbol: {collector.symbols}")
|
||||
logger.info(f"Collector data types: {[dt.value for dt in collector.data_types]}")
|
||||
|
||||
# Test 3: Create collector using convenience function
|
||||
logger.info("=== Testing Convenience Function ===")
|
||||
okx_collector = create_okx_collector(
|
||||
symbol='ETH-USDT',
|
||||
data_types=[DataType.TRADE],
|
||||
auto_restart=False
|
||||
)
|
||||
logger.info(f"Created OKX collector: {type(okx_collector).__name__}")
|
||||
logger.info(f"OKX collector symbol: {okx_collector.symbols}")
|
||||
|
||||
# Test 4: Create multiple collectors
|
||||
logger.info("=== Testing Multiple Collectors ===")
|
||||
configs = [
|
||||
ExchangeCollectorConfig('okx', 'BTC-USDT', [DataType.TRADE]),
|
||||
ExchangeCollectorConfig('okx', 'ETH-USDT', [DataType.ORDERBOOK]),
|
||||
ExchangeCollectorConfig('okx', 'SOL-USDT', [DataType.TRADE, DataType.ORDERBOOK])
|
||||
]
|
||||
|
||||
collectors = ExchangeFactory.create_multiple_collectors(configs)
|
||||
logger.info(f"Created {len(collectors)} collectors:")
|
||||
for i, collector in enumerate(collectors):
|
||||
logger.info(f" {i+1}. {type(collector).__name__} - {collector.symbols}")
|
||||
|
||||
# Test 5: Get exchange capabilities
|
||||
logger.info("=== Exchange Capabilities ===")
|
||||
okx_pairs = ExchangeFactory.get_supported_pairs('okx')
|
||||
okx_data_types = ExchangeFactory.get_supported_data_types('okx')
|
||||
logger.info(f"OKX supported pairs: {okx_pairs}")
|
||||
logger.info(f"OKX supported data types: {okx_data_types}")
|
||||
|
||||
logger.info("All factory tests completed successfully!")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Factory test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main test function."""
|
||||
logger = get_logger("main", verbose=True)
|
||||
logger.info("Testing exchange factory pattern...")
|
||||
|
||||
success = await test_factory_pattern()
|
||||
|
||||
if success:
|
||||
logger.info("Factory tests completed successfully!")
|
||||
else:
|
||||
logger.error("Factory tests 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)
|
||||
243
tests/test_okx_collector.py
Normal file
243
tests/test_okx_collector.py
Normal file
@@ -0,0 +1,243 @@
|
||||
#!/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)
|
||||
Reference in New Issue
Block a user