sientia-dataops-scouter_temporal/tests/activities/test_mongo.py

from datetime import datetime
from unittest.mock import ANY, MagicMock, patch

from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ

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+0000', DATETIME_FORMAT_MS_WITH_TZ
            ),
        }
    ]

    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+0000'},
    }


@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+0000', DATETIME_FORMAT_MS_WITH_TZ
            ),
        }
    ]

    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+0000',
        }
    )

    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+0000', DATETIME_FORMAT_MS_WITH_TZ
                )
            }
        },
        {'_id': 0},
    )

    assert result == {
        'name': {0: 'test1'},
        'value': {0: 1},
        'inserted_at': {0: '2023-01-01 12:00:00.000000+0000'},
    }


@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+0000',
            }
        )
    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,
        )