sientia-dataops-opc-ingestor/tests/unit/managers/test_data_manager.py

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