import asyncio import unittest from unittest.mock import AsyncMock, Mock, patch from datetime import datetime, timedelta, timezone from data.collector.collector_connection_manager import ConnectionManager from data.collector.collector_state_telemetry import CollectorStatus, CollectorStateAndTelemetry from data.common.data_types import DataType # Import for DataType enum class TestConnectionManager(unittest.IsolatedAsyncioTestCase): def setUp(self): self.mock_logger = Mock() self.mock_state_telemetry = AsyncMock(spec=CollectorStateAndTelemetry) self.mock_state_telemetry.logger = self.mock_logger self.mock_state_telemetry.update_status = Mock() self.mock_state_telemetry.set_should_be_running = Mock() self.mock_state_telemetry.set_connection_uptime_start = Mock() self.mock_state_telemetry._log_info = Mock() self.mock_state_telemetry._log_debug = Mock() self.mock_state_telemetry._log_warning = Mock() self.mock_state_telemetry._log_error = Mock() self.mock_state_telemetry._log_critical = Mock() self.manager = ConnectionManager( exchange_name="test_exchange", component_name="test_collector", logger=self.mock_logger, state_telemetry=self.mock_state_telemetry ) async def test_init(self): self.assertEqual(self.manager.exchange_name, "test_exchange") self.assertEqual(self.manager.component_name, "test_collector") self.assertEqual(self.manager._max_reconnect_attempts, 5) self.assertEqual(self.manager._reconnect_delay, 5.0) self.assertEqual(self.manager.logger, self.mock_logger) self.assertEqual(self.manager._state_telemetry, self.mock_state_telemetry) self.assertIsNone(self.manager._connection) self.assertEqual(self.manager._reconnect_attempts, 0) async def test_connect_success(self): mock_connect_logic = AsyncMock(return_value=True) result = await self.manager.connect(mock_connect_logic) self.assertTrue(result) mock_connect_logic.assert_called_once() self.assertTrue(self.manager._connection) self.mock_state_telemetry.set_connection_uptime_start.assert_called_once() self.mock_state_telemetry._log_info.assert_any_call("test_collector: Successfully connected to test_exchange") async def test_connect_failure(self): mock_connect_logic = AsyncMock(return_value=False) result = await self.manager.connect(mock_connect_logic) self.assertFalse(result) mock_connect_logic.assert_called_once() self.assertIsNone(self.manager._connection) self.mock_state_telemetry._log_error.assert_any_call("test_collector: Failed to connect to test_exchange", exc_info=False) async def test_connect_exception(self): mock_connect_logic = AsyncMock(side_effect=Exception("Connection error")) result = await self.manager.connect(mock_connect_logic) self.assertFalse(result) mock_connect_logic.assert_called_once() self.assertIsNone(self.manager._connection) self.mock_state_telemetry._log_error.assert_called_once_with("test_collector: Error during connection to test_exchange: Connection error", exc_info=True) async def test_disconnect_success(self): self.manager._connection = True # Simulate active connection mock_disconnect_logic = AsyncMock() await self.manager.disconnect(mock_disconnect_logic) mock_disconnect_logic.assert_called_once() self.assertIsNone(self.manager._connection) self.mock_state_telemetry._log_info.assert_any_call("test_collector: Disconnected from test_exchange") async def test_disconnect_no_active_connection(self): mock_disconnect_logic = AsyncMock() await self.manager.disconnect(mock_disconnect_logic) mock_disconnect_logic.assert_not_called() self.assertIsNone(self.manager._connection) self.mock_state_telemetry._log_info.assert_any_call("test_collector: Disconnecting from test_exchange data source") async def test_disconnect_exception(self): self.manager._connection = True mock_disconnect_logic = AsyncMock(side_effect=Exception("Disconnect error")) await self.manager.disconnect(mock_disconnect_logic) mock_disconnect_logic.assert_called_once() self.assertTrue(self.manager._connection) # Connection state remains unchanged on exception self.mock_state_telemetry._log_error.assert_called_once_with("test_collector: Error during disconnection from test_exchange: Disconnect error", exc_info=True) @patch('asyncio.sleep', new_callable=AsyncMock) async def test_handle_connection_error_reconnect_success(self, mock_sleep): mock_connect_logic = AsyncMock(return_value=True) mock_subscribe_logic = AsyncMock(return_value=True) self.manager._reconnect_attempts = 0 result = await self.manager.handle_connection_error( mock_connect_logic, mock_subscribe_logic, symbols=["BTC/USDT"], data_types=[DataType.CANDLE] ) self.assertTrue(result) self.assertEqual(self.manager._reconnect_attempts, 0) mock_connect_logic.assert_called_once() mock_subscribe_logic.assert_called_once_with(["BTC/USDT"], [DataType.CANDLE]) self.mock_state_telemetry.update_status.assert_any_call(CollectorStatus.RECONNECTING) self.mock_state_telemetry.update_status.assert_any_call(CollectorStatus.RUNNING) self.mock_state_telemetry._log_info.assert_any_call("test_collector: Reconnection successful for test_exchange") mock_sleep.assert_called_once_with(self.manager._reconnect_delay) @patch('asyncio.sleep', new_callable=AsyncMock) async def test_handle_connection_error_max_attempts_exceeded(self, mock_sleep): mock_connect_logic = AsyncMock(return_value=True) mock_subscribe_logic = AsyncMock(return_value=True) self.manager._reconnect_attempts = self.manager._max_reconnect_attempts result = await self.manager.handle_connection_error( mock_connect_logic, mock_subscribe_logic, symbols=["BTC/USDT"], data_types=[DataType.CANDLE] ) self.assertFalse(result) self.assertEqual(self.manager._reconnect_attempts, self.manager._max_reconnect_attempts + 1) mock_connect_logic.assert_not_called() mock_subscribe_logic.assert_not_called() self.mock_state_telemetry.update_status.assert_any_call(CollectorStatus.ERROR) self.mock_state_telemetry.set_should_be_running.assert_called_once_with(False) self.mock_state_telemetry._log_error.assert_any_call(f"test_collector: Max reconnection attempts ({self.manager._max_reconnect_attempts}) exceeded for test_exchange", exc_info=False) mock_sleep.assert_not_called() @patch('asyncio.sleep', new_callable=AsyncMock) async def test_handle_connection_error_connect_fails(self, mock_sleep): mock_connect_logic = AsyncMock(return_value=False) mock_subscribe_logic = AsyncMock(return_value=True) self.manager._reconnect_attempts = 0 result = await self.manager.handle_connection_error( mock_connect_logic, mock_subscribe_logic, symbols=["BTC/USDT"], data_types=[DataType.CANDLE] ) self.assertFalse(result) self.assertEqual(self.manager._reconnect_attempts, 1) mock_connect_logic.assert_called_once() mock_subscribe_logic.assert_not_called() self.mock_state_telemetry.update_status.assert_any_call(CollectorStatus.RECONNECTING) self.mock_state_telemetry._log_error.assert_called_once_with("test_collector: Failed to connect to test_exchange", exc_info=False) mock_sleep.assert_called_once_with(self.manager._reconnect_delay) @patch('asyncio.sleep', new_callable=AsyncMock) async def test_handle_connection_error_subscribe_fails(self, mock_sleep): mock_connect_logic = AsyncMock(return_value=True) mock_subscribe_logic = AsyncMock(return_value=False) self.manager._reconnect_attempts = 0 result = await self.manager.handle_connection_error( mock_connect_logic, mock_subscribe_logic, symbols=["BTC/USDT"], data_types=[DataType.CANDLE] ) self.assertFalse(result) self.assertEqual(self.manager._reconnect_attempts, 1) mock_connect_logic.assert_called_once() mock_subscribe_logic.assert_called_once_with(["BTC/USDT"], [DataType.CANDLE]) self.mock_state_telemetry.update_status.assert_any_call(CollectorStatus.RECONNECTING) mock_sleep.assert_called_once_with(self.manager._reconnect_delay) @patch('asyncio.sleep', new_callable=AsyncMock) async def test_handle_connection_error_exception_during_reconnect(self, mock_sleep): mock_connect_logic = AsyncMock(side_effect=Exception("Reconnect error")) mock_subscribe_logic = AsyncMock(return_value=True) self.manager._reconnect_attempts = 0 result = await self.manager.handle_connection_error( mock_connect_logic, mock_subscribe_logic, symbols=["BTC/USDT"], data_types=[DataType.CANDLE] ) self.assertFalse(result) self.assertEqual(self.manager._reconnect_attempts, 1) mock_connect_logic.assert_called_once() mock_subscribe_logic.assert_not_called() self.mock_state_telemetry.update_status.assert_any_call(CollectorStatus.RECONNECTING) self.mock_state_telemetry._log_error.assert_called_once_with("test_collector: Error during connection to test_exchange: Reconnect error", exc_info=True) mock_sleep.assert_called_once_with(self.manager._reconnect_delay)