109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
from pytest import fixture, mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
from ingestor.managers.data_manager import DataManager
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
'pod_id': 'localhost',
|
|
},
|
|
}
|
|
|
|
|
|
@fixture
|
|
@patch('ingestor.managers.data_manager.MongoDBRepository')
|
|
def data_manager(mongodb_repository):
|
|
data_manager = DataManager(
|
|
mongo_connection_string='mongodb://localhost:27017',
|
|
mongo_database='sientia',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metadata=metadata['metadata'],
|
|
metrics_controller=MagicMock(),
|
|
)
|
|
|
|
data_manager.send_notification = MagicMock()
|
|
data_manager.send_notification_async = AsyncMock()
|
|
data_manager.emit_metric = AsyncMock()
|
|
return data_manager
|
|
|
|
|
|
@patch('ingestor.managers.data_manager.MongoDBRepository')
|
|
def test___init___success(mongodb_repository):
|
|
logger_mock = MagicMock()
|
|
|
|
data_manager = DataManager(
|
|
metadata=metadata['metadata'],
|
|
mongo_connection_string='mongodb://localhost:27017',
|
|
mongo_database='sientia',
|
|
logger=logger_mock,
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(),
|
|
)
|
|
|
|
logger_mock.info.assert_any_call(
|
|
'DataManager initialized with MongoDB servers: mongodb://localhost:27017'
|
|
)
|
|
logger_mock.error.assert_not_called()
|
|
assert data_manager.connection_string == 'mongodb://localhost:27017'
|
|
|
|
|
|
def test_shutdown(data_manager):
|
|
close_mock = MagicMock()
|
|
data_manager.mongo_repository.close = close_mock
|
|
|
|
data_manager.shutdown()
|
|
close_mock.assert_called_once()
|
|
|
|
|
|
def test_shutdown_exception_mongo(data_manager):
|
|
data_manager.mongo_repository.close = MagicMock(side_effect=Exception('Test error'))
|
|
|
|
data_manager.shutdown()
|
|
|
|
data_manager.logger.error.assert_called_once_with('Error closing MongoDB client: Test error')
|
|
|
|
|
|
def test___del__(data_manager):
|
|
data_manager.shutdown = MagicMock()
|
|
data_manager.__del__()
|
|
data_manager.shutdown.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_publish(data_manager):
|
|
topic = 'test_topic'
|
|
data = {'key': 'value', 'name': 'test_tag'}
|
|
|
|
data_manager.mongo_repository.insert = AsyncMock()
|
|
|
|
await data_manager.publish(topic, data)
|
|
|
|
data_manager.mongo_repository.insert.assert_called_once_with(
|
|
collection_name=topic,
|
|
document={**data, 'inserted_at': ANY},
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_publish_error_mongo(data_manager):
|
|
data_manager.mongo_repository.insert = AsyncMock(side_effect=Exception('Test error'))
|
|
|
|
await data_manager.publish('test_topic', {'key': 'value'})
|
|
|
|
data_manager.send_notification_async.assert_called_once_with(
|
|
notification_id='MONGO_PRODUCER_ERROR_test_topic',
|
|
message='Error inserting message to MongoDB: Test error',
|
|
block='mongo_producer',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
metadata=metadata['metadata'],
|
|
)
|