from datetime import datetime from unittest.mock import ANY, MagicMock, call, patch from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from scouter.activities.mongodb import MongoDB, clear_mongo_id def test_clear_mongo_id(): """Test clear_mongo_id""" data = [ {'_id': '1', 'name': 'test1'}, {'_id': '2', 'name': [{ '_id': '3', 'name': 'test3' }]}, {'_id': '4', 'name': { '_id': '5', 'name': 'test2' }}, [{'_id': '6', 'name': 'test2'}] ] result = clear_mongo_id(data) assert result == [ {'name': 'test1'}, {'name': [{'name': 'test3'}]}, {'name': {'name': 'test2'}}, [{'name': 'test2'}] ] @patch('scouter.activities.mongodb.MongoClient') def test_mongodb___init__(mock_mongo_client): """Test MongoDB __init__""" mongo = MongoDB( connection_string='mongodb://localhost:27017', database_name='test_db', logger=MagicMock(), notification_handler=MagicMock() ) mock_mongo_client.assert_called_once_with( 'mongodb://localhost:27017', serverSelectionTimeoutMS=5000 ) mock_mongo_client.return_value.server_info.assert_called_once() mock_mongo_client.return_value.__getitem__.assert_called_once_with( 'test_db') assert mongo.client is not None assert mongo.database is not None @fixture @patch('scouter.activities.mongodb.MongoClient') def mongodb_activity(mock_mongo_client): """Test MongoDB activity""" mongo = MongoDB( connection_string='mongodb://localhost:27017', database_name='test_db', logger=MagicMock(), notification_handler=MagicMock() ) return mongo def test_shutdown_success(mongodb_activity): """Test shutdown""" mongodb_activity.shutdown() mongodb_activity.client.close.assert_called_once() def test_shutdown_error(mongodb_activity): """Test shutdown""" mongodb_activity.client.close = MagicMock(side_effect=Exception('test')) mongodb_activity.shutdown() mongodb_activity.client.close.assert_called_once() @mark.asyncio async def test_load_latest_data_none_last_data_timestamp(mongodb_activity): """Test load_latest_data""" collection = MagicMock() mongodb_activity.database.__getitem__.return_value = collection collection.find.return_value = [ { 'name': 'test1', 'value': 1, 'inserted_at': datetime.strptime( '2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f') } ] result = await mongodb_activity.load_latest_data({ 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, 'collection_name': 'test_collection', 'last_data_timestamp': None }) mongodb_activity.database.__getitem__.assert_called_once_with( 'test_collection') collection.find.assert_called_once_with( {}, {"_id": 0} ) assert result == { 'name': { 0: 'test1' }, 'value': { 0: 1 }, 'inserted_at': { 0: '2023-01-01 12:00:00.000000' } } @mark.asyncio async def test_load_latest_data_not_none_last_data_timestamp(mongodb_activity): """Test load_latest_data""" collection = MagicMock() mongodb_activity.database.__getitem__.return_value = collection collection.find.return_value = [ { 'name': 'test1', 'value': 1, 'inserted_at': datetime.strptime( '2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f') } ] result = await mongodb_activity.load_latest_data({ 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, 'collection_name': 'test_collection', 'last_data_timestamp': '2023-01-01 12:00:00.000000' }) mongodb_activity.database.__getitem__.assert_called_once_with( 'test_collection') collection.find.assert_called_once_with( { 'inserted_at': { '$gt': datetime.strptime( '2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f') } }, {"_id": 0} ) assert result == { 'name': { 0: 'test1' }, 'value': { 0: 1 }, 'inserted_at': { 0: '2023-01-01 12:00:00.000000' } } @mark.asyncio async def test_load_latest_data_error(mongodb_activity): """Test load_latest_data""" collection = MagicMock() mongodb_activity.send_notification = MagicMock() mongodb_activity.database.__getitem__.return_value = collection collection.find.side_effect = Exception('test') try: await mongodb_activity.load_latest_data({ 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, 'collection_name': 'test_collection', 'last_data_timestamp': '2023-01-01 12:00:00.000000' }) except Exception as e: assert str(e) == 'test' mongodb_activity.send_notification.assert_called_once_with( metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, notification_id='MONGO_LOAD_ERROR', message='Error loading data from MongoDB: test', block='load_latest_data', level=NotificationLevel.ERROR, attachment_content=ANY )