""" 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 from data.collector_types import 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() 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"])