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


@patch('scouter.activities.mongodb.MongoDBRepository')
def test_mongodb___init__(mock_mongodb_repository):
    """Test MongoDB __init__"""

    logger = MagicMock()
    notification_handler = MagicMock()

    mongo = MongoDB(
        connection_string='mongodb://localhost:27017',
        database_name='test_db',
        logger=logger,
        notification_handler=notification_handler,
    )

    mock_mongodb_repository.assert_called_once_with(
        connection_string='mongodb://localhost:27017',
        database_name='test_db',
        logger=logger,
        notification_handler=notification_handler,
    )

    assert mongo.mongodb_repository is not None


@fixture
@patch('scouter.activities.mongodb.MongoDBRepository')
def mongodb_activity(mock_mongodb_repository):
    """Test MongoDB activity"""

    mongo = MongoDB(
        connection_string='mongodb://localhost:27017',
        database_name='test_db',
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    mongo.logger = MagicMock()
    mongo.notification_handler = MagicMock()
    mongo.pod_id = 'localhost'
    return mongo


def test_close(mongodb_activity):
    """Test close"""
    mongodb_activity.close()

    mongodb_activity.mongodb_repository.close.assert_called_once()


def test_del(mongodb_activity):
    mongodb_activity.close = MagicMock()

    mongodb_activity.__del__()

    mongodb_activity.close.assert_called_once()


@mark.asyncio
async def test_load_latest_data_none_last_data_timestamp(mongodb_activity):
    """Test load_latest_data"""

    mongodb_activity.mongodb_repository.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.mongodb_repository.find.assert_called_once_with(
        collection_name='test_collection',
        filters={},
        metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
    )

    assert result == [
        {
            'name': 'test1',
            'value': 1,
            'inserted_at': '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"""

    mongodb_activity.mongodb_repository.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.mongodb_repository.find.assert_called_once_with(
        collection_name='test_collection',
        filters={
            'inserted_at': {
                '$gt': datetime.strptime(
                    '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
                )
            }
        },
        metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
    )

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


@mark.asyncio
async def test_load_latest_data_error(mongodb_activity):
    """Test load_latest_data"""
    mongodb_activity.mongodb_repository.find.side_effect = Exception('test')
    mongodb_activity.send_notification = MagicMock()

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