sientia-dataops-laborious-legacy/tests/laborious/activities/test_prediction_metadata.py

from unittest.mock import MagicMock, call, patch

from pytest import fixture
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring

from laborious.activities.prediction_metadata import PredictionMetadata


@fixture
def mongo() -> MagicMock:
    return MagicMock()


@fixture
def prediction_metadata(mongo: MagicMock):
    activity = PredictionMetadata(
        mongo_repository=mongo,
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(spec=MetricsController),
    )
    activity.error = MagicMock()
    return activity


metadata = {
    'metadata': {
        'model_id': 'test_model',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
        'schema_name': 'test_schedule',
    },
}


def test_write_prediction_metadata_no_documents_skips_mongo(prediction_metadata, mongo: MagicMock):
    # Act
    result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': []})

    # Assert
    assert result is None
    mongo.insert.assert_not_called()


def test_write_prediction_metadata_missing_documents_key_skips_mongo(
    prediction_metadata, mongo: MagicMock
):
    # Act
    result = prediction_metadata.write_prediction_metadata({**metadata})

    # Assert
    assert result is None
    mongo.insert.assert_not_called()


def test_write_prediction_metadata_documents_with_no_filters_skips_mongo(
    prediction_metadata, mongo: MagicMock
):
    # Arrange: a falsy document (None) and one with an empty `filters` list — neither contributes
    # anything to write.
    documents = [None, {'filters': []}]

    # Act
    result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': documents})

    # Assert
    assert result is None
    mongo.insert.assert_not_called()


def test_write_prediction_metadata_not_configured_returns_none():
    # Arrange
    activity = PredictionMetadata(
        mongo_repository=None,
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(spec=MetricsController),
    )
    activity.error = MagicMock()
    document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}

    # Act
    result = activity.write_prediction_metadata({**metadata, 'documents': [document]})

    # Assert
    assert result is None
    activity.error.assert_not_called()


def test_write_prediction_metadata_success_generates_id_and_inserts(
    prediction_metadata, mongo: MagicMock
):
    # Arrange
    document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 3}]}

    # Act
    result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': [document]})

    # Assert: a plain 24-hex-char string id (matches predictions.metadata_id's varchar(24)),
    # not a BSON ObjectId — `insert(...)` never hands one back, so it's generated client-side.
    assert result is not None
    assert len(result) == 24
    int(result, 16)  # raises ValueError if not hex

    collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
    assert collection == PredictionMetadata._PREDICTION_METADATA_COLLECTION
    assert inserted_document == {'filters': document['filters'], '_id': result}
    assert inserted_metadata == metadata['metadata']


def test_write_prediction_metadata_combines_multiple_documents(
    prediction_metadata, mongo: MagicMock
):
    # Arrange: two documents (e.g. from two different sources contributing metadata for the same
    # prediction cycle) — their `filters` entries are flattened into one combined document.
    documents = [
        {'filters': [{'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0}]},
        {
            'filters': [
                {'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}
            ]
        },
    ]

    # Act
    result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': documents})

    # Assert
    assert result is not None
    _, inserted_document, _ = mongo.insert.call_args.args
    assert inserted_document['filters'] == [
        {'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0},
        {'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2},
    ]


def test_write_prediction_metadata_insert_failure_swallowed(prediction_metadata, mongo: MagicMock):
    # Arrange
    mongo.insert.side_effect = RuntimeError('connection reset')
    document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}

    # Act
    result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': [document]})

    # Assert: the prediction write must never be blocked or failed by this
    assert result is None
    prediction_metadata.error.assert_called_once()


@patch.object(SientiaMonitoring, 'shutdown')
def test_close_disposes_mongo_repository(mock_shutdown):
    """``PredictionMetadata.close`` should close the Mongo client and clear the reference."""
    mongo = MagicMock()
    activity = PredictionMetadata(
        mongo_repository=mongo,
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(spec=MetricsController),
    )
    activity.close()
    mongo.close.assert_called_once()
    assert activity.mongo_repository is None
    # Other instances collected by GC during the test also hit the patched shutdown.
    assert mock_shutdown.call_args_list.count(call(activity)) == 1


@patch.object(SientiaMonitoring, 'shutdown')
def test_close_without_mongo_repository(mock_shutdown):
    """When no Mongo repository is configured, ``close`` only shuts down monitoring."""
    activity = PredictionMetadata(
        mongo_repository=None,
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(spec=MetricsController),
    )
    activity.close()
    # Other instances collected by GC during the test also hit the patched shutdown.
    assert mock_shutdown.call_args_list.count(call(activity)) == 1