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

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


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

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


@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
        )