sientia-dataops-model-manager/tests/activities/test_activities.py

"""Unit tests for Activities orchestrator (constructor, shutdown, destructor)."""

from unittest.mock import MagicMock, Mock, patch

import pytest
from sientia_do.observability.sientia_monitoring import SientiaMonitoring

from model_manager.activities.activities import Activities


def _postgres():
    return {
        'host': 'h',
        'port': 5432,
        'user': 'u',
        'password': 'p',
        'dbname': 'db',
        'min_connections': 1,
        'max_connections': 2,
    }


def _mlflow():
    return {'url': 'http://mlflow:5000', 'username': 'u', 'password': 'p'}


def _minio(endpoint_url: str):
    return {
        'endpoint_url': endpoint_url,
        'access_key': 'a',
        'secret_key': 's',
        'region': 'r',
        'use_ssl': True,
        'default_bucket': 'test-bucket',
    }


@pytest.mark.parametrize(
    'endpoint,expected_endpoint',
    [
        ('http://minio:9000', 'minio:9000'),
        ('https://minio:9000', 'minio:9000'),
        ('minio:9000', 'minio:9000'),
    ],
)
def test_activities_strips_minio_endpoint_scheme(endpoint, expected_endpoint):
    with (
        patch(
            'model_manager.activities.activities.ExperimentTracking.__init__',
            Mock(return_value=None),
        ),
        patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.SientiaMLflowRepository') as m_mlflow,
        patch('model_manager.activities.activities.MinioRepository') as m_minio,
    ):
        Activities(
            postgres_config=_postgres(),
            mlflow_config=_mlflow(),
            minio_config=_minio(endpoint),
            plugin_store=MagicMock(),
            logger=MagicMock(),
            notification_handler=MagicMock(),
            metrics_controller=MagicMock(),
        )
    m_minio.assert_called_once()
    assert m_minio.call_args.kwargs['endpoint'] == expected_endpoint
    assert m_minio.call_args.kwargs['bucket'] == 'test-bucket'
    m_mlflow.assert_called_once()


def test_activities_shutdown_calls_parents():
    with (
        patch(
            'model_manager.activities.activities.ExperimentTracking.__init__',
            Mock(return_value=None),
        ),
        patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.SientiaMLflowRepository'),
        patch('model_manager.activities.activities.MinioRepository'),
        patch('model_manager.activities.activities.ExperimentTracking.close') as m_close,
        patch('model_manager.activities.activities.SientiaMonitoring.shutdown') as m_mon,
    ):
        a = Activities(
            postgres_config=_postgres(),
            mlflow_config=_mlflow(),
            minio_config=_minio('http://x:9000'),
            plugin_store=MagicMock(),
            logger=MagicMock(),
            notification_handler=MagicMock(),
            metrics_controller=MagicMock(),
        )
        with patch.object(SientiaMonitoring, 'info', Mock()):
            a.shutdown()
    m_close.assert_called_once()
    m_mon.assert_called_once()


def test_activities_del_with_engine_runs_without_error():
    with (
        patch(
            'model_manager.activities.activities.ExperimentTracking.__init__',
            Mock(return_value=None),
        ),
        patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.SientiaMLflowRepository'),
        patch('model_manager.activities.activities.MinioRepository'),
    ):
        a = Activities(
            postgres_config=_postgres(),
            mlflow_config=_mlflow(),
            minio_config=_minio('http://x:9000'),
            plugin_store=MagicMock(),
            logger=MagicMock(),
            notification_handler=MagicMock(),
            metrics_controller=MagicMock(),
        )
    a.engine = MagicMock()
    Activities.__del__(a)


def test_activities_del_without_engine_runs_without_error():
    with (
        patch(
            'model_manager.activities.activities.ExperimentTracking.__init__',
            Mock(return_value=None),
        ),
        patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
        patch('model_manager.activities.activities.SientiaMLflowRepository'),
        patch('model_manager.activities.activities.MinioRepository'),
    ):
        a = Activities(
            postgres_config=_postgres(),
            mlflow_config=_mlflow(),
            minio_config=_minio('http://x:9000'),
            plugin_store=MagicMock(),
            logger=MagicMock(),
            notification_handler=MagicMock(),
            metrics_controller=MagicMock(),
        )
    Activities.__del__(a)