from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from ingestor.managers.data_manager import DataManager, redact_uri 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' @patch('ingestor.managers.data_manager.MongoDBRepository') def test___init___does_not_log_the_credentials(mongodb_repository): logger_mock = MagicMock() uri = 'mongodb://root:s3cr3t@mongodb.mongodb.svc.cluster.local:27017' data_manager = DataManager( metadata=metadata['metadata'], mongo_connection_string=uri, mongo_database='sientia', logger=logger_mock, notification_handler=MagicMock(), metrics_controller=MagicMock(), ) logged = ' '.join(str(c) for c in logger_mock.info.call_args_list) assert 's3cr3t' not in logged assert 'root' not in logged logger_mock.info.assert_any_call( 'DataManager initialized with MongoDB servers: ' 'mongodb://***@mongodb.mongodb.svc.cluster.local:27017' ) assert data_manager.connection_string == uri @pytest.mark.parametrize( ('uri', 'expected'), [ ('mongodb://localhost:27017', 'mongodb://localhost:27017'), ( 'mongodb://u:p@h1:27017,h2:27017/db?authSource=admin', 'mongodb://***@h1:27017,h2:27017/db?authSource=admin', ), ('mongodb://u:p@h:27017?authSource=admin', 'mongodb://***@h:27017?authSource=admin'), ('mongodb+srv://u@cluster.example.com', 'mongodb+srv://***@cluster.example.com'), ('u:p@h:27017', '***@h:27017'), ('mongodb://h:27017/db?x=a@b', 'mongodb://h:27017/db?x=a@b'), ], ) def test_redact_uri(uri, expected): assert redact_uri(uri) == expected 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'], )