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

from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from sientia_do.temporal.activities.postgres import Postgres
from laborious.activities.activities import Activities
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.opc import OPC


@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):

    postgres_config = {
        'host': 'localhost',
        'port': 5432,
        'user': 'postgres',
        'password': 'postgres',
        'dbname': 'postgres',
        'min_connections': 1,
        'max_connections': 10
    }

    mlflow_config = {
        'host': 'localhost',
        'port': 5000,
        'username': 'mlflow',
        'password': 'mlflow'
    }

    opc_config = {
        'bootstrap_servers': 'localhost:9092',
        'polling_time': 1000,
        'group_id': 'test-group'
    }

    logger = MagicMock()
    notification_handler = MagicMock()

    activities = Activities(
        postgres_config=postgres_config,
        mlflow_config=mlflow_config,
        opc_config=opc_config,
        logger=logger,
        notification_handler=notification_handler
    )

    assert isinstance(activities, Activities)
    assert isinstance(activities, Postgres)
    assert isinstance(activities, MLFlow)
    assert isinstance(activities, OPC)
    assert isinstance(activities, Gates)

    mock_postgres_init.assert_called_once_with(
        ANY,
        host=postgres_config['host'],
        port=postgres_config['port'],
        user=postgres_config['user'],
        password=postgres_config['password'],
        dbname=postgres_config['dbname'],
        min_connections=postgres_config['min_connections'],
        max_connections=postgres_config['max_connections'],
        logger=logger,
        notification_handler=notification_handler
    )

    mock_mlflow_init.assert_called_once_with(
        ANY,
        mlflow_host=mlflow_config['host'],
        mlflow_port=mlflow_config['port'],
        mlflow_username=mlflow_config['username'],
        mlflow_password=mlflow_config['password'],
        logger=logger,
        notification_handler=notification_handler
    )

    mock_opc_init.assert_called_once_with(
        ANY,
        opc_servers=opc_config,
        logger=logger,
        notification_handler=notification_handler
    )

    mock_gates_init.assert_called_once_with(
        ANY,
        logger=logger,
        notification_handler=notification_handler
    )


@mark.asyncio
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
async def test_prepare_activity(_mock_opc_init,
                                _mock_mlflow_init, _mock_postgres_init):
    postgres_config = {
        'host': 'localhost',
        'port': 5432,
        'user': 'postgres',
        'password': 'postgres',
        'dbname': 'postgres',
        'min_connections': 1,
        'max_connections': 10
    }

    mlflow_config = {
        'host': 'localhost',
        'port': 5000,
        'username': 'mlflow',
        'password': 'mlflow'
    }

    opc_config = {
        'bootstrap_servers': 'localhost:9092',
        'polling_time': 1000,
        'group_id': 'test-group'
    }

    logger = MagicMock()
    notification_handler = MagicMock()

    activities = Activities(
        postgres_config=postgres_config,
        mlflow_config=mlflow_config,
        opc_config=opc_config,
        logger=logger,
        notification_handler=notification_handler
    )

    input_data = {
        'workflow_name': 'test-workflow-name',
        'schedule_name': 'test-schedule-name',
        'model_name': 'test-model-name',
        'model_id': 'test-model-id'
    }

    await activities.prepare_activity(input_data)

    assert activities.notification_handler.base_notification.pipeline == input_data[
        'workflow_name']
    assert activities.notification_handler.base_notification.trigger == input_data[
        'schedule_name']
    assert activities.notification_handler.base_notification.model_name == input_data[
        'model_name']
    assert activities.notification_handler.base_notification.model_id == input_data[
        'model_id']


@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())
def test_shutdown(mock_opc_init,
                  _mock_mlflow_init, mock_postgres_init):
    postgres_config = {
        'host': 'localhost',
        'port': 5432,
        'user': 'postgres',
        'password': 'postgres',
        'dbname': 'postgres',
        'min_connections': 1,
        'max_connections': 10
    }

    mlflow_config = {
        'host': 'localhost',
        'port': 5000,
        'username': 'mlflow',
        'password': 'mlflow'
    }

    opc_config = {
        'bootstrap_servers': 'localhost:9092',
        'polling_time': 1000,
        'group_id': 'test-group'
    }

    logger = MagicMock()
    notification_handler = MagicMock()

    activities = Activities(
        postgres_config=postgres_config,
        mlflow_config=mlflow_config,
        opc_config=opc_config,
        logger=logger,
        notification_handler=notification_handler
    )

    activities.shutdown()
    mock_opc_init.shutdown.assert_called_once()
    mock_postgres_init.close.assert_called_once()