- Extracted callback management logic into a new `CallbackDispatcher` class, promoting separation of concerns and enhancing modularity. - Updated `BaseDataCollector` to utilize the `CallbackDispatcher` for adding, removing, and notifying data callbacks, improving code clarity and maintainability. - Refactored related methods to ensure consistent error handling and logging practices. - Added unit tests for the `CallbackDispatcher` to validate its functionality and ensure robust error handling. These changes streamline the callback management architecture, aligning with project standards for maintainability and performance.
112 lines
5.1 KiB
Python
112 lines
5.1 KiB
Python
import asyncio
|
|
import unittest
|
|
from unittest.mock import AsyncMock, Mock
|
|
from datetime import datetime
|
|
|
|
from data.collector.collector_callback_dispatcher import CallbackDispatcher
|
|
from data.common.data_types import DataType, MarketDataPoint
|
|
|
|
|
|
class TestCallbackDispatcher(unittest.IsolatedAsyncioTestCase):
|
|
|
|
def setUp(self):
|
|
self.mock_logger = Mock()
|
|
self.dispatcher = CallbackDispatcher(
|
|
component_name="test_dispatcher",
|
|
logger=self.mock_logger
|
|
)
|
|
|
|
async def test_init(self):
|
|
self.assertEqual(self.dispatcher.component_name, "test_dispatcher")
|
|
self.assertEqual(self.dispatcher.logger, self.mock_logger)
|
|
self.assertIsInstance(self.dispatcher._data_callbacks, dict)
|
|
self.assertGreater(len(self.dispatcher._data_callbacks), 0) # Ensure all DataType enums are initialized
|
|
|
|
async def test_add_data_callback(self):
|
|
mock_callback = Mock()
|
|
data_type = DataType.CANDLE
|
|
|
|
self.dispatcher.add_data_callback(data_type, mock_callback)
|
|
self.assertIn(mock_callback, self.dispatcher._data_callbacks[data_type])
|
|
self.mock_logger.debug.assert_called_with(f"test_dispatcher: Added callback for {data_type.value} data")
|
|
|
|
# Test adding same callback twice (should not add)
|
|
self.dispatcher.add_data_callback(data_type, mock_callback)
|
|
self.assertEqual(self.dispatcher._data_callbacks[data_type].count(mock_callback), 1)
|
|
|
|
async def test_remove_data_callback(self):
|
|
mock_callback = Mock()
|
|
data_type = DataType.TRADE
|
|
|
|
self.dispatcher.add_data_callback(data_type, mock_callback)
|
|
self.assertIn(mock_callback, self.dispatcher._data_callbacks[data_type])
|
|
|
|
self.dispatcher.remove_data_callback(data_type, mock_callback)
|
|
self.assertNotIn(mock_callback, self.dispatcher._data_callbacks[data_type])
|
|
self.mock_logger.debug.assert_called_with(f"test_dispatcher: Removed callback for {data_type.value} data")
|
|
|
|
# Test removing non-existent callback (should do nothing)
|
|
self.dispatcher.remove_data_callback(data_type, Mock())
|
|
# No error should be raised and log should not be called again for removal
|
|
|
|
async def test_notify_callbacks_sync(self):
|
|
mock_sync_callback = Mock()
|
|
data_type = DataType.TICKER
|
|
data_point = MarketDataPoint("exchange", "symbol", datetime.now(), data_type, {"price": 100})
|
|
|
|
self.dispatcher.add_data_callback(data_type, mock_sync_callback)
|
|
await self.dispatcher.notify_callbacks(data_point)
|
|
|
|
mock_sync_callback.assert_called_once_with(data_point)
|
|
|
|
async def test_notify_callbacks_async(self):
|
|
mock_async_callback = AsyncMock()
|
|
data_type = DataType.ORDERBOOK
|
|
data_point = MarketDataPoint("exchange", "symbol", datetime.now(), data_type, {"bids": [], "asks": []})
|
|
|
|
self.dispatcher.add_data_callback(data_type, mock_async_callback)
|
|
await self.dispatcher.notify_callbacks(data_point)
|
|
|
|
mock_async_callback.assert_called_once_with(data_point)
|
|
|
|
async def test_notify_callbacks_mixed(self):
|
|
mock_sync_callback = Mock()
|
|
mock_async_callback = AsyncMock()
|
|
data_type = DataType.BALANCE
|
|
data_point = MarketDataPoint("exchange", "symbol", datetime.now(), data_type, {"asset": "BTC", "balance": 0.5})
|
|
|
|
self.dispatcher.add_data_callback(data_type, mock_sync_callback)
|
|
self.dispatcher.add_data_callback(data_type, mock_async_callback)
|
|
await self.dispatcher.notify_callbacks(data_point)
|
|
|
|
mock_sync_callback.assert_called_once_with(data_point)
|
|
mock_async_callback.assert_called_once_with(data_point)
|
|
|
|
async def test_notify_callbacks_exception_handling(self):
|
|
def failing_sync_callback(data): raise ValueError("Sync error")
|
|
async def failing_async_callback(data): raise TypeError("Async error")
|
|
|
|
mock_successful_callback = Mock()
|
|
|
|
data_type = DataType.CANDLE
|
|
data_point = MarketDataPoint("exchange", "symbol", datetime.now(), data_type, {})
|
|
|
|
self.dispatcher.add_data_callback(data_type, failing_sync_callback)
|
|
self.dispatcher.add_data_callback(data_type, failing_async_callback)
|
|
self.dispatcher.add_data_callback(data_type, mock_successful_callback)
|
|
|
|
await self.dispatcher.notify_callbacks(data_point)
|
|
|
|
mock_successful_callback.assert_called_once_with(data_point)
|
|
self.assertEqual(self.mock_logger.error.call_count, 2)
|
|
self.mock_logger.error.assert_any_call(f"test_dispatcher: Error in data callback for {data_type.value} {data_point.symbol}: Sync error", exc_info=True)
|
|
self.mock_logger.error.assert_any_call(f"test_dispatcher: Error in data callback for {data_type.value} {data_point.symbol}: Async error", exc_info=True)
|
|
|
|
async def test_notify_callbacks_no_callbacks(self):
|
|
data_type = DataType.TICKER
|
|
data_point = MarketDataPoint("exchange", "symbol", datetime.now(), data_type, {})
|
|
|
|
# No callbacks added
|
|
await self.dispatcher.notify_callbacks(data_point)
|
|
self.mock_logger.error.assert_not_called() # No errors should be logged
|
|
self.mock_logger.debug.assert_not_called() # No debug logs from notify |