33f2110f19
- Extracted `OHLCVData` and validation logic into a new `common/ohlcv_data.py` module, promoting better organization and reusability. - Updated `BaseDataCollector` to utilize the new `validate_ohlcv_data` function for improved data validation, enhancing code clarity and maintainability. - Refactored imports in `data/__init__.py` to reflect the new structure, ensuring consistent access to common data types and exceptions. - Removed redundant data validation logic from `BaseDataCollector`, streamlining its responsibilities. - Added unit tests for `OHLCVData` and validation functions to ensure correctness and reliability. These changes improve the architecture of the data module, aligning with project standards for maintainability and performance.
353 lines
12 KiB
Python
353 lines
12 KiB
Python
"""
|
|
Unit tests for the CollectorManager class.
|
|
"""
|
|
|
|
import asyncio
|
|
import pytest
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from utils.logger import get_logger
|
|
from data.collector_manager import CollectorManager, ManagerStatus, CollectorConfig
|
|
from data.base_collector import BaseDataCollector, DataType, CollectorStatus
|
|
|
|
|
|
class MockDataCollector(BaseDataCollector):
|
|
"""Mock implementation of BaseDataCollector for testing."""
|
|
|
|
def __init__(self, exchange_name: str, symbols: list, auto_restart: bool = True):
|
|
super().__init__(exchange_name, symbols, [DataType.TICKER], auto_restart=auto_restart)
|
|
self.connected = False
|
|
self.subscribed = False
|
|
self.should_fail_connect = False
|
|
self.should_fail_subscribe = False
|
|
self.fail_count = 0
|
|
|
|
async def _actual_connect(self) -> bool:
|
|
"""Implementation of actual connection logic for testing."""
|
|
if self.should_fail_connect and self.fail_count < 2:
|
|
self.fail_count += 1
|
|
return False
|
|
await asyncio.sleep(0.01)
|
|
self.connected = True
|
|
return True
|
|
|
|
async def _actual_disconnect(self) -> None:
|
|
"""Implementation of actual disconnection logic for testing."""
|
|
await asyncio.sleep(0.01)
|
|
self.connected = False
|
|
self.subscribed = False
|
|
|
|
async def connect(self) -> bool:
|
|
"""Connect using the connection manager."""
|
|
return await self._connection_manager.connect(self._actual_connect)
|
|
|
|
async def disconnect(self) -> None:
|
|
"""Disconnect using the connection manager."""
|
|
await self._connection_manager.disconnect(self._actual_disconnect)
|
|
|
|
async def subscribe_to_data(self, symbols: list, data_types: list) -> bool:
|
|
if self.should_fail_subscribe:
|
|
return False
|
|
if not self.connected:
|
|
return False
|
|
self.subscribed = True
|
|
return True
|
|
|
|
async def unsubscribe_from_data(self, symbols: list, data_types: list) -> bool:
|
|
self.subscribed = False
|
|
return True
|
|
|
|
async def _process_message(self, message) -> None:
|
|
# No message processing in mock
|
|
pass
|
|
|
|
async def _handle_messages(self) -> None:
|
|
# Simulate light processing
|
|
await asyncio.sleep(0.1)
|
|
|
|
|
|
class TestCollectorManager:
|
|
"""Test cases for CollectorManager."""
|
|
|
|
@pytest.fixture
|
|
def manager(self):
|
|
"""Create a test manager instance."""
|
|
test_logger = get_logger("test_manager_logger")
|
|
return CollectorManager("test_manager", global_health_check_interval=1.0, logger=test_logger)
|
|
|
|
@pytest.fixture
|
|
def mock_collector(self):
|
|
"""Create a mock collector."""
|
|
return MockDataCollector("okx", ["BTC-USDT", "ETH-USDT"])
|
|
|
|
def test_initialization(self, manager):
|
|
"""Test manager initialization."""
|
|
assert manager.manager_name == "test_manager"
|
|
assert manager.status == ManagerStatus.STOPPED
|
|
assert len(manager._collectors) == 0
|
|
assert len(manager._enabled_collectors) == 0
|
|
|
|
def test_add_collector(self, manager, mock_collector):
|
|
"""Test adding a collector to the manager."""
|
|
# Add collector
|
|
manager.add_collector(mock_collector)
|
|
|
|
assert len(manager._collectors) == 1
|
|
assert len(manager._enabled_collectors) == 1
|
|
|
|
# Verify collector is in the collections
|
|
collector_names = manager.list_collectors()
|
|
assert len(collector_names) == 1
|
|
assert collector_names[0].startswith("okx_")
|
|
|
|
# Test with custom config using a different collector instance
|
|
mock_collector2 = MockDataCollector("binance", ["ETH-USDT"])
|
|
config = CollectorConfig(
|
|
name="custom_collector",
|
|
exchange="binance",
|
|
symbols=["ETH-USDT"],
|
|
data_types=["ticker"],
|
|
enabled=False
|
|
)
|
|
manager.add_collector(mock_collector2, config)
|
|
assert len(manager._collectors) == 2
|
|
assert len(manager._enabled_collectors) == 1 # Still 1 since second is disabled
|
|
|
|
def test_remove_collector(self, manager, mock_collector):
|
|
"""Test removing a collector from the manager."""
|
|
# Add then remove
|
|
manager.add_collector(mock_collector)
|
|
collector_names = manager.list_collectors()
|
|
collector_name = collector_names[0]
|
|
|
|
success = manager.remove_collector(collector_name)
|
|
assert success
|
|
assert len(manager._collectors) == 0
|
|
assert len(manager._enabled_collectors) == 0
|
|
|
|
# Test removing non-existent collector
|
|
success = manager.remove_collector("non_existent")
|
|
assert not success
|
|
|
|
def test_enable_disable_collector(self, manager, mock_collector):
|
|
"""Test enabling and disabling collectors."""
|
|
manager.add_collector(mock_collector)
|
|
collector_name = manager.list_collectors()[0]
|
|
|
|
# Initially enabled
|
|
assert collector_name in manager._enabled_collectors
|
|
|
|
# Disable
|
|
success = manager.disable_collector(collector_name)
|
|
assert success
|
|
assert collector_name not in manager._enabled_collectors
|
|
|
|
# Enable again
|
|
success = manager.enable_collector(collector_name)
|
|
assert success
|
|
assert collector_name in manager._enabled_collectors
|
|
|
|
# Test with non-existent collector
|
|
success = manager.enable_collector("non_existent")
|
|
assert not success
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_stop_manager(self, manager, mock_collector):
|
|
"""Test starting and stopping the manager."""
|
|
# Add a collector
|
|
manager.add_collector(mock_collector)
|
|
|
|
# Start manager
|
|
success = await manager.start()
|
|
assert success
|
|
assert manager.status == ManagerStatus.RUNNING
|
|
|
|
# Wait a bit for collectors to start
|
|
await asyncio.sleep(0.2)
|
|
|
|
# Check collector is running
|
|
running_collectors = manager.get_running_collectors()
|
|
assert len(running_collectors) == 1
|
|
|
|
# Stop manager
|
|
await manager.stop()
|
|
assert manager.status == ManagerStatus.STOPPED
|
|
|
|
# Check collector is stopped
|
|
running_collectors = manager.get_running_collectors()
|
|
assert len(running_collectors) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_collector(self, manager, mock_collector):
|
|
"""Test restarting a specific collector."""
|
|
manager.add_collector(mock_collector)
|
|
await manager.start()
|
|
|
|
collector_name = manager.list_collectors()[0]
|
|
|
|
# Wait for collector to start
|
|
await asyncio.sleep(0.2)
|
|
|
|
# Restart the collector
|
|
success = await manager.restart_collector(collector_name)
|
|
assert success
|
|
|
|
# Check statistics
|
|
status = manager.get_status()
|
|
assert status['statistics']['restarts_performed'] >= 1
|
|
|
|
await manager.stop()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_monitoring(self, manager):
|
|
"""Test health monitoring and auto-restart functionality."""
|
|
# Create a collector that will fail initially
|
|
failing_collector = MockDataCollector("test", ["BTC-USDT"], auto_restart=True)
|
|
failing_collector.should_fail_connect = True
|
|
|
|
manager.add_collector(failing_collector)
|
|
await manager.start()
|
|
|
|
# Wait for health checks
|
|
await asyncio.sleep(2.5) # More than health check interval
|
|
|
|
# Check that restarts were attempted
|
|
status = manager.get_status()
|
|
failed_collectors = manager.get_failed_collectors()
|
|
|
|
# The collector should have been marked as failed and restart attempts made
|
|
assert len(failed_collectors) >= 0 # May have recovered
|
|
|
|
await manager.stop()
|
|
|
|
def test_get_status(self, manager, mock_collector):
|
|
"""Test status reporting."""
|
|
manager.add_collector(mock_collector)
|
|
|
|
status = manager.get_status()
|
|
|
|
assert status['manager_status'] == 'stopped'
|
|
assert status['total_collectors'] == 1
|
|
assert len(status['enabled_collectors']) == 1
|
|
assert 'statistics' in status
|
|
assert 'collectors' in status
|
|
|
|
def test_get_collector_status(self, manager, mock_collector):
|
|
"""Test getting individual collector status."""
|
|
manager.add_collector(mock_collector)
|
|
collector_name = manager.list_collectors()[0]
|
|
|
|
collector_status = manager.get_collector_status(collector_name)
|
|
|
|
assert collector_status is not None
|
|
assert collector_status['name'] == collector_name
|
|
assert 'config' in collector_status
|
|
assert 'status' in collector_status
|
|
assert 'health' in collector_status
|
|
|
|
# Test non-existent collector
|
|
non_existent_status = manager.get_collector_status("non_existent")
|
|
assert non_existent_status is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_all_collectors(self, manager):
|
|
"""Test restarting all collectors."""
|
|
# Add multiple collectors
|
|
collector1 = MockDataCollector("okx", ["BTC-USDT"])
|
|
collector2 = MockDataCollector("binance", ["ETH-USDT"])
|
|
|
|
manager.add_collector(collector1)
|
|
manager.add_collector(collector2)
|
|
|
|
await manager.start()
|
|
await asyncio.sleep(0.2) # Let them start
|
|
|
|
# Restart all
|
|
results = await manager.restart_all_collectors()
|
|
|
|
assert len(results) == 2
|
|
assert all(success for success in results.values())
|
|
|
|
await manager.stop()
|
|
|
|
def test_get_running_and_failed_collectors(self, manager, mock_collector):
|
|
"""Test getting running and failed collector lists."""
|
|
manager.add_collector(mock_collector)
|
|
|
|
# Initially no running collectors
|
|
running = manager.get_running_collectors()
|
|
failed = manager.get_failed_collectors()
|
|
|
|
assert len(running) == 0
|
|
# Note: failed might be empty since collector hasn't started yet
|
|
|
|
def test_collector_config(self):
|
|
"""Test CollectorConfig dataclass."""
|
|
config = CollectorConfig(
|
|
name="test_collector",
|
|
exchange="okx",
|
|
symbols=["BTC-USDT", "ETH-USDT"],
|
|
data_types=["ticker", "trade"],
|
|
auto_restart=True,
|
|
health_check_interval=30.0,
|
|
enabled=True
|
|
)
|
|
|
|
assert config.name == "test_collector"
|
|
assert config.exchange == "okx"
|
|
assert len(config.symbols) == 2
|
|
assert len(config.data_types) == 2
|
|
assert config.auto_restart is True
|
|
assert config.enabled is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_with_connection_failures():
|
|
"""Test manager handling collectors with connection failures."""
|
|
manager = CollectorManager("test_manager", global_health_check_interval=0.5)
|
|
|
|
# Create a collector that fails connection initially
|
|
failing_collector = MockDataCollector("failing_exchange", ["BTC-USDT"])
|
|
failing_collector.should_fail_connect = True
|
|
|
|
manager.add_collector(failing_collector)
|
|
|
|
# Start manager
|
|
success = await manager.start()
|
|
assert success # Manager should start even if collectors fail
|
|
|
|
# Wait for some health checks
|
|
await asyncio.sleep(1.5)
|
|
|
|
# Check that the failing collector is detected
|
|
failed_collectors = manager.get_failed_collectors()
|
|
status = manager.get_status()
|
|
|
|
# The collector should be in failed state or have restart attempts
|
|
assert status['statistics']['restarts_performed'] >= 0
|
|
|
|
await manager.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_graceful_shutdown():
|
|
"""Test that manager shuts down gracefully even with problematic collectors."""
|
|
manager = CollectorManager("test_manager")
|
|
|
|
# Add multiple collectors
|
|
for i in range(3):
|
|
collector = MockDataCollector(f"exchange_{i}", ["BTC-USDT"])
|
|
manager.add_collector(collector)
|
|
|
|
await manager.start()
|
|
await asyncio.sleep(0.2)
|
|
|
|
# Stop should complete even if collectors take time
|
|
await manager.stop()
|
|
|
|
assert manager.status == ManagerStatus.STOPPED
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"]) |