Code import - branch 0.6.0

tests/laborious/activities/test_activities.py

@@ -0,0 +1,229 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
+from pytest import mark
+
+from laborious.activities.activities import Activities
+from laborious.activities.api import API
+from laborious.activities.gates import Gates
+from laborious.activities.mlflow import MLFlow
+from laborious.activities.model_metrics import ModelMetrics
+from laborious.activities.opc import OPC
+from laborious.activities.storage import Storage
+
+
+@patch('laborious.activities.activities.Storage.__init__')
+@patch('laborious.activities.activities.MLFlow.__init__')
+@patch('laborious.activities.activities.OPC.__init__')
+@patch('laborious.activities.activities.Gates.__init__')
+@patch('laborious.activities.activities.ModelMetrics.__init__')
+@patch('laborious.activities.activities.API.__init__')
+@patch('laborious.activities.activities.MinioRepository')
+@patch('laborious.activities.activities.MetricsController')
+def test___init__(
+    mock_metrics_controller,
+    mock_minio_repository,
+    mock_api_init,
+    mock_model_metrics_init,
+    mock_gates_init,
+    mock_opc_init,
+    mock_mlflow_init,
+    mock_storage_init,
+):
+    postgres_config = {
+        'host': 'localhost',
+        'port': 5432,
+        'user': 'postgres',
+        'password': 'postgres',
+        'dbname': 'postgres',
+        'min_connections': 1,
+        'max_connections': 10,
+    }
+
+    minio_config = {
+        'endpoint_url': 'localhost:9000',
+        'access_key': 'minio',
+        'secret_key': 'minio123',
+        'default_bucket': 'test',
+        'retention_hours': 24,
+        'secure': False,
+    }
+
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
+
+    opc_config = {
+        'bootstrap_servers': 'localhost:9092',
+        'polling_time': 1000,
+        'group_id': 'test-group',
+    }
+
+    pi_web_api_config = {
+        'base_url': 'https://test-pi-server.com',
+        'auth_type': 'bearer',
+        'auth_token': 'test_token',
+    }
+
+    logger = MagicMock()
+    notification_handler = MagicMock()
+
+    activities = Activities(
+        postgres_config=postgres_config,
+        mlflow_config=mlflow_config,
+        minio_config=minio_config,
+        opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
+        logger=logger,
+        notification_handler=notification_handler,
+    )
+
+    assert isinstance(activities, Activities)
+    assert isinstance(activities, Storage)
+    assert isinstance(activities, MLFlow)
+    assert isinstance(activities, OPC)
+    assert isinstance(activities, Gates)
+    assert isinstance(activities, ModelMetrics)
+    assert isinstance(activities, API)
+
+    mock_storage_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'],
+        retention_hours=minio_config['retention_hours'],
+        minio_repository=mock_minio_repository.return_value,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    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'],
+        minio_repository=mock_minio_repository.return_value,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_opc_init.assert_called_once_with(
+        ANY,
+        opc_servers=opc_config,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_gates_init.assert_called_once_with(
+        ANY,
+        minio_repository=mock_minio_repository.return_value,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_model_metrics_init.assert_called_once_with(
+        ANY,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_api_init.assert_called_once_with(
+        ANY,
+        base_url=pi_web_api_config['base_url'],
+        auth_type=pi_web_api_config['auth_type'],
+        auth_token=pi_web_api_config['auth_token'],
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_minio_repository.assert_called_once_with(
+        endpoint=minio_config['endpoint_url'],
+        access_key=minio_config['access_key'],
+        secret_key=minio_config['secret_key'],
+        bucket=minio_config['default_bucket'],
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+        secure=minio_config['secure'],
+    )
+
+
+@mark.asyncio
+@patch('laborious.activities.activities.Storage')
+@patch('laborious.activities.activities.MLFlow')
+@patch('laborious.activities.activities.OPC')
+@patch('laborious.activities.activities.Gates')
+@patch('laborious.activities.activities.ModelMetrics')
+@patch('laborious.activities.activities.API')
+@patch('laborious.activities.activities.MinioRepository')
+async def test_shutdown(
+    _mock_minio_repository,
+    mock_api_init,
+    mock_model_metrics_init,
+    mock_gates_init,
+    mock_opc_init,
+    mock_mlflow_init,
+    mock_storage_init,
+):
+    mock_opc_init.close = AsyncMock()
+    postgres_config = {
+        'host': 'localhost',
+        'port': 5432,
+        'user': 'postgres',
+        'password': 'postgres',
+        'dbname': 'postgres',
+        'min_connections': 1,
+        'max_connections': 10,
+    }
+
+    minio_config = {
+        'endpoint_url': 'localhost:9000',
+        'access_key': 'minio',
+        'secret_key': 'minio123',
+        'default_bucket': 'test',
+        'retention_hours': 24,
+        'secure': False,
+    }
+
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
+
+    opc_config = {
+        'bootstrap_servers': 'localhost:9092',
+        'polling_time': 1000,
+        'group_id': 'test-group',
+    }
+
+    pi_web_api_config = {
+        'base_url': 'https://test-pi-server.com',
+        'auth_type': 'bearer',
+        'auth_token': 'test_token',
+    }
+
+    logger = MagicMock()
+    notification_handler = MagicMock()
+
+    activities = Activities(
+        postgres_config=postgres_config,
+        mlflow_config=mlflow_config,
+        minio_config=minio_config,
+        opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
+        logger=logger,
+        notification_handler=notification_handler,
+    )
+
+    await activities.shutdown()
+    mock_opc_init.close.assert_called_once()
+    mock_storage_init.close.assert_called_once()
+    mock_mlflow_init.close.assert_called_once()
+    mock_gates_init.close.assert_called_once()
+    mock_model_metrics_init.close.assert_called_once()
+    mock_api_init.close.assert_called_once()

tests/laborious/activities/test_api.py

@@ -0,0 +1,492 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+import pytest_asyncio
+from pytest import fixture, mark
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.activities.api import API, PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+def _create_mock_dataframe(to_dict_return=None):
+    """Helper function to create a mocked DataFrame for testing."""
+    mock_df = MagicMock()
+    mock_head = MagicMock()
+
+    def get_column_values(key):
+        if key == 'prediction':
+            return MagicMock(values=[0.75])
+        elif key == 'prediction_confidence':
+            return MagicMock(values=[0.95])
+        else:
+            return MagicMock(values=['2024-01-01T00:00:00+00:00'])
+
+    mock_head.__getitem__.side_effect = get_column_values
+    mock_df.head.return_value = mock_head
+
+    if to_dict_return is None:
+        to_dict_return = {
+            'prediction': [0.75],
+            'prediction_confidence': [0.95],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        }
+    mock_df.to_dict.return_value = to_dict_return
+
+    return mock_df
+
+
+@fixture
+def base_input_data():
+    """Base input data for PI Web API tests."""
+    return {
+        **metadata,
+        'data': {
+            'prediction': [0.75],
+            'prediction_confidence': [0.95],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        },
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com/piwebapi',
+            'prediction_tags': {'tag1': 'web_id_1'},
+            'confidence_tags': {'tag2': 'web_id_2'},
+        },
+    }
+
+
+@patch('laborious.activities.api.PIWebAPIClient')
+def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client):
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
+    api_instance = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    with patch.object(
+        SientiaMonitoring,
+        'get_core_labels',
+        return_value={
+            'pod_id': 'test_pod',
+            'runtime': 'local',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+            'operation_type': 'write_pi_web_api_data',
+        },
+    ):
+        labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata'])
+    assert labels['operation_type'] == 'write_pi_web_api_data'
+    assert labels == {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'operation_type': 'write_pi_web_api_data',
+    }
+
+
+@patch('laborious.activities.api.PIWebAPIClient')
+def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
+    api_instance = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    with patch.object(
+        SientiaMonitoring,
+        'get_core_labels',
+        return_value={
+            'pod_id': 'test_pod',
+            'runtime': 'k8s',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+            'operation_type': 'write',
+        },
+    ):
+        labels = api_instance.get_pi_web_api_core_labels(
+            metadata=metadata['metadata'], operation_type='write'
+        )
+    assert labels['operation_type'] == 'write'
+    assert labels['runtime'] == 'k8s'
+
+
+def test__init__():
+    api = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+
+    assert api.pi_web_api_client is not None
+
+
+@pytest_asyncio.fixture
+@patch('laborious.activities.api.PIWebAPIClient')
+def api(mock_pi_web_api_client):
+    mock_client = MagicMock()
+    mock_client.write_value = AsyncMock()
+    mock_client.close = MagicMock()
+    mock_client.base_url = 'https://test-pi-server.com'
+    mock_pi_web_api_client.return_value = mock_client
+
+    api_instance = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    api_instance.send_notification_async = AsyncMock()
+    api_instance.info = MagicMock()
+    api_instance.error = MagicMock()
+    api_instance.emit_metric = AsyncMock()
+    api_instance.get_core_labels = MagicMock(
+        return_value={
+            'pod_id': 'test_pod',
+            'runtime': 'local',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+        }
+    )
+    return api_instance
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_data):
+    input_data = {
+        **base_input_data,
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com/piwebapi',
+            'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
+            'confidence_tags': {'tag3': 'web_id_3', 'tag4': 'web_id_4'},
+        },
+    }
+
+    mock_dataframe.return_value = _create_mock_dataframe()
+
+    # Mock successful responses
+    api.pi_web_api_client.write_value.side_effect = [
+        [{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}],
+        [{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}],
+    ]
+
+    result = await api.write_pi_web_api_data(input_data)
+
+    api.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=['web_id_1', 'web_id_2'],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.75,
+                },
+                metadata=metadata['metadata'],
+            ),
+            call(
+                web_ids=['web_id_3', 'web_id_4'],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.95,
+                },
+                metadata=metadata['metadata'],
+            ),
+        ]
+    )
+
+    assert result == {
+        'prediction': [0.75],
+        'prediction_confidence': [0.95],
+        'timestamp': ['2024-01-01T00:00:00+00:00'],
+    }
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
+    mock_dataframe.return_value = _create_mock_dataframe(
+        {
+            'prediction': [0.75],
+            'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        }
+    )
+
+    api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')
+
+    result = await api.write_pi_web_api_data(base_input_data)
+
+    api.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
+        message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
+        block='write_pi_web_api_data',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+    assert result['prediction_confidence'][0] == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert api.pi_web_api_client.write_value.call_count == 1
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_input_data):
+    mock_dataframe.return_value = _create_mock_dataframe()
+
+    # First call succeeds, second fails
+    api.pi_web_api_client.write_value.side_effect = [
+        [{'WebId': 'web_id_1', 'Errors': []}],
+        Exception('Confidence write failed'),
+    ]
+
+    result = await api.write_pi_web_api_data(base_input_data)
+
+    api.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
+        message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
+        block='write_pi_web_api_data',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+    assert result == {
+        'prediction': [0.75],
+        'prediction_confidence': [0.95],
+        'timestamp': ['2024-01-01T00:00:00+00:00'],
+    }
+    assert api.pi_web_api_client.write_value.call_count == 2
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
+    input_data = {
+        **base_input_data,
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com/piwebapi',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+    }
+
+    mock_dataframe.return_value = _create_mock_dataframe()
+
+    # Mock empty responses
+    api.pi_web_api_client.write_value.side_effect = [
+        [],
+        [],
+    ]
+
+    result = await api.write_pi_web_api_data(input_data)
+
+    api.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=[],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.75,
+                },
+                metadata=metadata['metadata'],
+            ),
+            call(
+                web_ids=[],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.95,
+                },
+                metadata=metadata['metadata'],
+            ),
+        ]
+    )
+
+    assert result == {
+        'prediction': [0.75],
+        'prediction_confidence': [0.95],
+        'timestamp': ['2024-01-01T00:00:00+00:00'],
+    }
+
+
+@mark.asyncio
+async def test_close(api):
+    api.close()
+
+    api.pi_web_api_client.close.assert_called_once()
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_success(api):
+    """Test successful processing of PI Web API response with all tags written."""
+    response_data = [
+        {'WebId': 'web_id_1', 'Errors': []},
+        {'WebId': 'web_id_2', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == 0
+    assert message == ''
+    assert api.emit_metric.call_count == 2
+    # Verify that emit_metric was called with correct tags structure
+    call_args_list = api.emit_metric.call_args_list
+    assert len(call_args_list) == 2
+    # Check that all calls include core_labels and tag_name
+    for call_args in call_args_list:
+        assert 'tag_name' in call_args.kwargs['tags']
+        assert call_args.kwargs['tags']['tag_name'] in ['tag1', 'tag2']
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_with_errors(api):
+    """Test processing response with errors in some tags."""
+    response_data = [
+        {'WebId': 'web_id_1', 'Errors': ['Error writing tag']},
+        {'WebId': 'web_id_2', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
+    )
+    assert api.emit_metric.call_count == 2
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_missing_tags(api):
+    """Test processing response when number of written tags doesn't match expected."""
+    response_data = [
+        {'WebId': 'web_id_1', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
+    )
+    api.send_notification_async.assert_called_once()
+    call_args = api.send_notification_async.call_args
+    assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
+    assert call_args.kwargs['level'] == NotificationLevel.ERROR
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_missing_webid(api):
+    """Test processing response when WebId is missing in response item."""
+    response_data = [
+        {'Errors': []},
+        {'WebId': 'web_id_2', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
+    )
+    api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_missing_tag_name(api):
+    """Test processing response when tag name is not found for WebId."""
+    response_data = [
+        {'WebId': 'unknown_web_id', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
+    )
+    api.error.assert_any_call(
+        'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
+    )

tests/laborious/activities/test_mlflow.py

@@ -0,0 +1,629 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+import numpy as np
+from pytest import fixture, mark, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.mlflow import MLFlow
+
+
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
+@patch('laborious.activities.mlflow.MLFlowRepository')
+@patch('laborious.activities.mlflow.MinioRepository')
+def test___init__(mock_minio_repository, mock_mlflow_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    minio_repo = mock_minio_repository(
+        endpoint='localhost:9000',
+        access_key='minio',
+        secret_key='minio123',
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+        bucket='test',
+    )
+    mlflow = MLFlow(
+        mlflow_host='http://localhost',
+        mlflow_port=5000,
+        mlflow_username='admin',
+        mlflow_password='admin',
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    assert mlflow.mlflow_host == 'http://localhost'
+    assert mlflow.mlflow_port == 5000
+    assert mlflow.mlflow_username == 'admin'
+    assert mlflow.mlflow_password == 'admin'
+
+    mock_mlflow_repository.assert_called_once_with(
+        'http://localhost:5000', 'admin', 'admin', ANY, ANY, ANY
+    )
+
+    mock_minio_repository.assert_called_once_with(
+        endpoint='localhost:9000',
+        access_key='minio',
+        secret_key='minio123',
+        logger=ANY,
+        notification_handler=ANY,
+        metrics_controller=ANY,
+        bucket='test',
+    )
+
+
+@fixture
+@patch('laborious.activities.mlflow.MLFlowRepository')
+@patch('laborious.activities.mlflow.MinioRepository')
+def mlflow(mock_minio_repository, mock_mlflow_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    minio_repo = mock_minio_repository(
+        endpoint='localhost:9000',
+        access_key='minio',
+        secret_key='minio123',
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+        bucket='test',
+    )
+    mlflow = MLFlow(
+        mlflow_host='http://localhost:5000',
+        mlflow_port=5000,
+        mlflow_username='admin',
+        mlflow_password='admin',
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    mlflow.model_monitoring_repository = AsyncMock()
+    mlflow.minio_repository = AsyncMock()
+
+    mlflow.send_notification = MagicMock()
+    mlflow.emit_metric = AsyncMock()
+    mlflow.send_notification_async = AsyncMock()
+    mlflow.error = MagicMock()
+    mlflow.debug = MagicMock()
+    mlflow.info = MagicMock()
+    mlflow.warning = MagicMock()
+    mlflow.critical = MagicMock()
+
+    return mlflow
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+async def test_request_transform_success(mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    transform_response = {'success': True, 'content': MagicMock()}
+    mlflow.model_monitoring_repository.transform.return_value = transform_response
+
+    data_mock.sort_values.return_value = data_mock
+    data_mock.drop_duplicates.return_value = data_mock
+    data_mock.pivot.return_value = data_mock
+
+    response_data = await mlflow.request_transform(input_data)
+
+    mlflow.model_monitoring_repository.transform.assert_called_once_with(
+        'test_model', data_mock, {}, metadata['metadata']
+    )
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
+
+
+@mark.asyncio
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+async def test_request_transform_failure(mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    transform_response = {'success': False, 'message': 'Transform failed'}
+    mlflow.model_monitoring_repository.transform.return_value = transform_response
+
+    data_mock.sort_values.return_value = data_mock
+    data_mock.drop_duplicates.return_value = data_mock
+    data_mock.pivot.return_value = data_mock
+
+    response_data = await mlflow.request_transform(input_data)
+
+    mock_from_dataframe.assert_called_once_with(
+        dataframe=None,
+        minio_repo=mlflow.minio_repository,
+        model_name='test_model',
+        operation='transform',
+        status=transform_response,
+        workflow_metadata=metadata['metadata'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+    assert response_data == mock_from_dataframe.return_value
+
+
+@mark.asyncio
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    predict_response = {'success': True, 'content': MagicMock()}
+    mlflow.model_monitoring_repository.predict.return_value = predict_response
+
+    response_data = await mlflow.request_predict(input_data)
+
+    data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
+    mock_to_datetime.assert_called_once_with(
+        data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
+    )
+    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
+
+    mlflow.model_monitoring_repository.predict.assert_called_once_with(
+        'test_model', data_mock, {}, metadata['metadata']
+    )
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
+
+
+@mark.asyncio
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    predict_response = {'success': False, 'message': 'Predict failed'}
+    mlflow.model_monitoring_repository.predict.return_value = predict_response
+
+    response_data = await mlflow.request_predict(input_data)
+
+    mock_from_dataframe.assert_called_once_with(
+        dataframe=None,
+        minio_repo=mlflow.minio_repository,
+        model_name='test_model',
+        operation='predict',
+        status=predict_response,
+        workflow_metadata=metadata['metadata'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+    assert response_data == mock_from_dataframe.return_value
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
+        'success': True,
+        'experiment': 'test_experiment',
+        'message': 'Model retrained successfully.',
+    }
+
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=raw_data)
+
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    timestamp = raw_data.__getitem__.return_value.max.return_value
+
+    raw_data.sort_values.assert_not_called()
+    raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
+    raw_data = raw_data.drop_duplicates.return_value
+
+    raw_data.drop.assert_has_calls(
+        [
+            call(columns=['model_id'], inplace=True, errors='ignore'),
+            call(columns=['created_at'], inplace=True, errors='ignore'),
+        ]
+    )
+    raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
+    raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
+
+    raw_data = raw_data.pivot.return_value
+
+    raw_data.__setitem__.assert_has_calls(
+        [
+            call('timestamp', raw_data.index),
+            call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
+            call('timestamp', mock_to_datetime.return_value),
+        ]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
+    )
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
+        data=raw_data,
+        model_name='test_model',
+        model_config={
+            'target': 'target',
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'pyfunc',
+        },
+        metadata=metadata['metadata'],
+    )
+
+    assert response == {
+        'success': True,
+        'experiment': 'test_experiment',
+        'message': 'Model retrained successfully.',
+        'timestamp': timestamp,
+    }
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow):
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
+        'success': True,
+        'experiment': 'test_experiment',
+        'message': 'Model retrained successfully.',
+    }
+
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=raw_data)
+
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    assert response['success'] is True
+    mlflow.minio_repository.download_file.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
+        'success': False,
+        'traceback': 'test_traceback',
+        'message': 'Model retrained failed.',
+    }
+
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=raw_data)
+
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    timestamp = raw_data.__getitem__.return_value.max.return_value
+
+    raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
+    raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
+        subset=['variable', 'timestamp'], keep='first'
+    )
+    raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
+
+    raw_data.drop.assert_has_calls(
+        [
+            call(columns=['model_id'], inplace=True, errors='ignore'),
+            call(columns=['created_at'], inplace=True, errors='ignore'),
+        ]
+    )
+    raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
+    raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
+
+    raw_data = raw_data.pivot.return_value
+
+    raw_data.__setitem__.assert_has_calls(
+        [
+            call('timestamp', raw_data.index),
+            call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
+            call('timestamp', mock_to_datetime.return_value),
+        ]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
+    )
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
+        data=raw_data,
+        model_name='test_model',
+        model_config={
+            'target': 'target',
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'pyfunc',
+        },
+        metadata=metadata['metadata'],
+    )
+
+    mlflow.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='RETRAIN_MODEL_ERROR',
+        message='Error retraining model test_model: Model retrained failed.',
+        block='retrain_model',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+    assert response == {
+        'success': False,
+        'traceback': 'test_traceback',
+        'message': 'Model retrained failed.',
+        'timestamp': timestamp,
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_data_error(mlflow):
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    assert response == {
+        'success': False,
+        'message': "Error loading retrain data: 'data'",
+        'traceback': ANY,
+        'timestamp': ANY,
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_data_error_no_minio_repository(mlflow):
+    mlflow.minio_repository = None
+
+    with raises(ValueError) as e:
+        await mlflow.retrain_model(
+            {
+                **metadata,
+                'object_key': 'test_object_key',
+                'model_name': 'test_model',
+                'model_config': {
+                    'target': 'target',
+                    'transform_flavor': 'sklearn',
+                    'predict_flavor': 'pyfunc',
+                },
+            }
+        )
+
+        assert str(e.value) == 'Minio repository not initialized'
+
+
+@mark.asyncio
+async def test_update_production_model(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 1,
+        'experiment': 'test',
+        'timestamp': 2,
+        'status': 'success',
+    }
+
+    response = await mlflow.update_production_model(input_data)
+
+    mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
+        experiment='test', model_name='test_model', metadata=metadata['metadata']
+    )
+
+    assert response == mlflow.model_monitoring_repository.update_production_model.return_value
+
+
+@mark.asyncio
+async def test_update_production_model_error(mlflow):
+    mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
+        'Error updating production model'
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 1,
+        'experiment': 'test',
+        'timestamp': 2,
+        'status': 'success',
+    }
+
+    try:
+        await mlflow.update_production_model(input_data)
+    except Exception as e:
+        assert str(e) == 'Error updating production model'
+        mlflow.send_notification_async.assert_called_once_with(
+            metadata=metadata['metadata'],
+            notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
+            message='Error updating production model test_model: Error updating production model',
+            block='update_production_model',
+            level=NotificationLevel.ERROR,
+            attachment_content=ANY,
+        )
+    else:
+        raise AssertionError('No exception raised')
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_get_reference_data_success(mock_to_datetime, mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    # Mock reference data DataFrame
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [
+        {'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
+        {'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
+    ]
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
+
+    # Act
+    result = await mlflow.get_reference_data(input_data)
+
+    # Assert
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )
+    mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
+
+    mock_reference_data.to_dict.assert_called_once_with(orient='records')
+    assert result == mock_reference_data.to_dict.return_value
+
+
+@mark.asyncio
+async def test_get_reference_data_not_found(mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
+
+    # Act
+    result = await mlflow.get_reference_data(input_data)
+
+    # Assert
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )
+    mlflow.warning.assert_called_once_with(
+        'Reference data not found for model test_model', metadata['metadata']
+    )
+    assert result is None
+
+
+@mark.asyncio
+async def test_get_reference_data_exception(mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
+        'Error loading artifact'
+    )
+
+    # Act & Assert
+    with raises(Exception) as e:
+        await mlflow.get_reference_data(input_data)
+
+    assert str(e.value) == 'Error loading artifact'
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )

tests/laborious/activities/test_opc.py

@@ -0,0 +1,738 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+import pytest_asyncio
+from pandas import DataFrame
+from pytest import mark
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.activities.opc import (
+    OPC,
+    OPC_COMMENT_SEPARATOR,
+    OPC_RECONNECT_IN_PROGRESS_COMMENT,
+    OPC_SESSION_BAD_COMMENT_PREFIX,
+    OPC_SESSION_BAD_CONFIDENCE,
+    OPC_WRITTING_ERROR_CONFIDENCE,
+    OPC_WRITTING_ERROR_MESSAGE,
+    _apply_opc_write_error,
+)
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+def test__init__():
+    servers = {'server1': {'id': 'server1'}}
+    opc = OPC(
+        opc_servers=servers,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+
+    assert opc.opc_servers == servers
+    assert opc.opc_repository == {}
+
+
+@mark.asyncio
+@patch('laborious.activities.opc.OpcRepository')
+@patch('laborious.activities.opc.OPC.send_notification_async')
+async def test_init_opc(mock_send_notification, mock_opc_repository):
+    mock_logger = MagicMock()
+    mock_metrics_controller = AsyncMock()
+    server1 = MagicMock(
+        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
+    )
+    server2 = MagicMock(
+        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
+    )
+    server3 = MagicMock(
+        connect=AsyncMock(
+            return_value=(
+                False,
+                {
+                    'notification_id': 'OPC_CONNECTION_ERROR_server3',
+                    'message': 'Failed to connect to OPC server: Test error',
+                    'block': 'opc_repository',
+                    'level': NotificationLevel.ERROR,
+                    'attachment_content': 'Test error',
+                },
+            )
+        ),
+        write_data=AsyncMock(return_value=(True, {})),
+    )
+    mock_opc_repository.side_effect = [server1, server2, server3]
+    mock_notification_handler = MagicMock()
+    servers = {
+        'server1': {
+            'server_name': 'server1',
+            'id': 'server1',
+            'url': 'http://localhost:8080',
+            'server_uri': 'opc.tcp://localhost:4840',
+            'cert_path': '',
+            'private_key_path': '',
+            'server_cert_path': '',
+            'reconnection_interval': 60,
+        },
+        'server2': {
+            'server_name': 'server2',
+            'id': 'server2',
+            'url': 'http://localhost:8080',
+            'server_uri': 'opc.tcp://localhost:4840',
+            'cert_path': '',
+            'private_key_path': '',
+            'server_cert_path': '',
+            'reconnection_interval': 60,
+        },
+        'server3': {
+            'server_name': 'server3',
+            'id': 'server3',
+            'url': 'http://localhost:8080',
+            'server_uri': 'opc.tcp://localhost:4840',
+            'cert_path': '',
+            'private_key_path': '',
+            'server_cert_path': '',
+            'reconnection_interval': 60,
+        },
+    }
+    opc = OPC(
+        opc_servers=servers,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+        metrics_controller=mock_metrics_controller,
+    )
+    await opc.init_opc()
+
+    assert opc.opc_servers == servers
+    assert opc.logger == mock_logger
+    assert opc.notification_handler == mock_notification_handler
+    assert opc.opc_repository['server1'] == server1
+    assert opc.opc_repository['server2'] == server2
+
+    mock_opc_repository.assert_has_calls(
+        [
+            call(
+                opc_id='server1',
+                server_name='server1',
+                url='http://localhost:8080',
+                logger=mock_logger,
+                server_uri='opc.tcp://localhost:4840',
+                cert_path='',
+                private_key_path='',
+                server_cert_path='',
+                notification_handler=mock_notification_handler,
+                reconnection_interval=60,
+                metrics_controller=mock_metrics_controller,
+            ),
+        ]
+    )
+    mock_opc_repository.assert_has_calls(
+        [
+            call(
+                opc_id='server2',
+                server_name='server2',
+                url='http://localhost:8080',
+                logger=mock_logger,
+                server_uri='opc.tcp://localhost:4840',
+                cert_path='',
+                private_key_path='',
+                server_cert_path='',
+                notification_handler=mock_notification_handler,
+                reconnection_interval=60,
+                metrics_controller=mock_metrics_controller,
+            )
+        ]
+    )
+
+    server1.connect.assert_called_once()
+    server2.connect.assert_called_once()
+
+    mock_send_notification.assert_has_calls(
+        [
+            call(
+                metadata={
+                    'model_id': '-',
+                    'model_name': '-',
+                    'workflow_name': '-',
+                    'schedule_name': 'INITIALIZATION',
+                },
+                notification_id='OPC_CONNECTION_ERROR_server3',
+                message='Failed to connect to OPC server: Test error',
+                block='opc_repository',
+                level=NotificationLevel.ERROR,
+                attachment_content=ANY,
+            )
+        ]
+    )
+
+
+@pytest_asyncio.fixture
+@patch('laborious.activities.opc.OpcRepository')
+async def opc(mock_opc_repository):
+    servers = {
+        'server1': {
+            'id': 'server1',
+            'server_name': 'server1',
+            'url': 'http://localhost:8080',
+            'server_uri': 'opc.tcp://localhost:4840',
+            'cert_path': '',
+            'private_key_path': '',
+            'server_cert_path': '',
+            'reconnection_interval': 60,
+        }
+    }
+
+    mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
+    mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
+    opc = OPC(
+        opc_servers=servers,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    await opc.init_opc()
+    opc.send_notification = MagicMock()
+    opc.send_notification_async = AsyncMock()
+    opc.emit_metric = AsyncMock()
+    return opc
+
+
+WRITE_DATA_CASES = [
+    ('tag1', 'int', 50),
+    ('tag2', 'float', 50.5),
+    ('tag3', 'bool', True),
+    ('tag4', 'string', 'test'),
+]
+
+
+@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
+@mark.asyncio
+async def test_write_data_success(opc, tag, data_type, data):
+    opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})
+
+    response_time, error_info = await opc.write_data(
+        server_id='server1',
+        tag=tag,
+        data=data,
+        data_type=data_type,
+        tag_type='prediction',
+        metadata=metadata,
+    )
+    assert response_time == 0.1
+    assert error_info is None
+    opc.opc_repository['server1'].write_data.assert_called_once_with(tag, data, data_type, metadata)
+
+
+@mark.asyncio
+async def test_write_data_failed(opc):
+    opc.opc_repository['server1'].write_data.return_value = (
+        False,
+        {
+            'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
+            'message': 'Failed to write data to OPC server: Test error',
+            'block': 'opc_repository',
+            'level': NotificationLevel.ERROR,
+            'attachment_content': 'Test error',
+        },
+    )
+
+    response_time, error_info = await opc.write_data(
+        server_id='server1',
+        tag='tag1',
+        data=50,
+        data_type='int',
+        tag_type='prediction',
+        metadata=metadata,
+    )
+    assert response_time is None
+    assert error_info is not None
+
+    opc.send_notification_async.assert_called_once_with(
+        metadata=metadata,
+        notification_id='OPC_WRITE_DATA_ERROR_server1',
+        message='Failed to write data to OPC server: Test error',
+        block='opc_repository',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+@mark.asyncio
+async def test_write_data_exception(opc):
+    opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
+
+    try:
+        await opc.write_data(
+            server_id='server1',
+            tag='tag1',
+            data=50,
+            data_type='int',
+            tag_type='prediction',
+            metadata=metadata,
+        )
+
+    except Exception:
+        opc.send_notification_async.assert_called_once_with(
+            metadata=metadata,
+            notification_id='WRITE_OPC_PREDICTION_ERROR',
+            message='Error writing data to OPC server: Test error',
+            block='write_opc_data',
+            level=NotificationLevel.ERROR,
+            attachment_content=ANY,
+        )
+
+    else:
+        raise AssertionError('Expected an exception to be raised')
+
+
+@mark.parametrize(
+    'error_info,initial_seen,initial_status,initial_reconnect,expected',
+    [
+        (None, False, None, False, (False, None, False)),
+        ({}, False, None, False, (False, None, False)),
+        (
+            {'opc_error_kind': 'session_bad', 'opc_status': 'BadSessionIdInvalid'},
+            False,
+            None,
+            False,
+            (True, 'BadSessionIdInvalid', False),
+        ),
+        (
+            {'opc_error_kind': 'session_bad', 'opc_status': 'NewStatus'},
+            True,
+            'OldStatus',
+            False,
+            (True, 'NewStatus', False),
+        ),
+        (
+            {'opc_error_kind': 'session_bad'},
+            True,
+            'KeptStatus',
+            False,
+            (True, 'KeptStatus', False),
+        ),
+        (
+            {'opc_error_kind': 'reconnect_in_progress'},
+            False,
+            None,
+            False,
+            (False, None, True),
+        ),
+        (
+            {'opc_error_kind': 'other'},
+            True,
+            'Status',
+            True,
+            (True, 'Status', True),
+        ),
+    ],
+)
+def test_apply_opc_write_error(
+    error_info, initial_seen, initial_status, initial_reconnect, expected
+):
+    result = _apply_opc_write_error(
+        error_info,
+        initial_seen,
+        initial_status,
+        initial_reconnect,
+    )
+    assert result == expected
+
+
+@mark.asyncio
+async def test_write_tags_from_config_prediction_success(opc):
+    opc.write_data = AsyncMock(return_value=(0.1, None))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    tags_config = {'tag1': {'data_type': 'float'}}
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config=tags_config,
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': 0.1}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+    opc.write_data.assert_called_once_with(
+        server_id='server1',
+        tag='tag1',
+        data=0.75,
+        data_type='float',
+        tag_type='prediction',
+        metadata=metadata['metadata'],
+    )
+
+
+@mark.asyncio
+async def test_write_tags_from_config_confidence_success(opc):
+    opc.write_data = AsyncMock(return_value=(0.2, None))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    tags_config = {'tag2': {'data_type': 'float'}}
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config=tags_config,
+        data=data,
+        data_column='prediction_confidence',
+        tag_type='confidence',
+        log_label='Confidence data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag2': 0.2}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+    opc.write_data.assert_called_once_with(
+        server_id='server1',
+        tag='tag2',
+        data=0.95,
+        data_type='float',
+        tag_type='confidence',
+        metadata=metadata['metadata'],
+    )
+
+
+@mark.asyncio
+async def test_write_tags_from_config_write_failure(opc):
+    opc.write_data = AsyncMock(return_value=(None, {}))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config={'tag1': {'data_type': 'float'}},
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': None}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+
+
+@mark.asyncio
+async def test_write_tags_from_config_session_bad(opc):
+    opc.write_data = AsyncMock(
+        return_value=(
+            None,
+            {
+                'opc_error_kind': 'session_bad',
+                'opc_status': 'BadSessionIdInvalid',
+            },
+        )
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config={'tag1': {'data_type': 'float'}},
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': None}
+    assert session_bad is True
+    assert opc_status == 'BadSessionIdInvalid'
+    assert reconnect is False
+
+
+@mark.asyncio
+async def test_write_tags_from_config_reconnect_in_progress(opc):
+    opc.write_data = AsyncMock(
+        return_value=(
+            None,
+            {'opc_error_kind': 'reconnect_in_progress'},
+        )
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config={'tag1': {'data_type': 'float'}},
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': None}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is True
+
+
+@mark.asyncio
+async def test_manage_output_tags_success(opc):
+    opc._write_tags_from_config = AsyncMock(
+        side_effect=[
+            ({'tag1': 0.1}, False, None, False),
+            ({'tag2': 0.1}, False, None, False),
+        ]
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {
+        'prediction_tags': {'tag1': {'data_type': 'float'}},
+        'confidence_tags': {'tag2': {'data_type': 'float'}},
+    }
+
+    output_data, opc_metrics, session_bad, opc_status, reconnect = await opc.manage_output_tags(
+        server_id='server1',
+        config=config,
+        data=data,
+        metadata=metadata['metadata'],
+    )
+
+    assert output_data is True
+    assert opc_metrics == {'tag1': 0.1, 'tag2': 0.1}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+    assert opc._write_tags_from_config.await_count == 2
+
+
+@mark.asyncio
+async def test_manage_output_tags_failed(opc):
+    opc._write_tags_from_config = AsyncMock(
+        side_effect=[
+            ({'tag1': 0.1}, False, None, False),
+            ({'tag2': None}, False, None, False),
+        ]
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {
+        'prediction_tags': {'tag1': {'data_type': 'float'}},
+        'confidence_tags': {'tag2': {'data_type': 'float'}},
+    }
+
+    output_data, opc_metrics, _, _, _ = await opc.manage_output_tags(
+        server_id='server1',
+        config=config,
+        data=data,
+        metadata=metadata['metadata'],
+    )
+
+    assert output_data is False
+    assert opc_metrics == {'tag1': 0.1, 'tag2': None}
+
+
+@mark.asyncio
+async def test_manage_output_tags_do_nothing(opc):
+    opc._write_tags_from_config = AsyncMock()
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {'_invalid_key': {'tag1': {'data_type': 'float'}}}
+
+    output_data, opc_metrics, _, _, _ = await opc.manage_output_tags(
+        server_id='server1',
+        config=config,
+        data=data,
+        metadata=metadata['metadata'],
+    )
+
+    assert output_data is True
+    assert opc_metrics == {}
+    opc._write_tags_from_config.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.opc.DataFrame')
+async def test_write_opc_data_success(mock_dataframe, opc):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
+        'opc_output_config': {
+            'server1': {
+                'prediction_tags': {'tag1': {'data_type': 'float'}},
+                'confidence_tags': {'tag2': {'data_type': 'float'}},
+            }
+        },
+    }
+
+    # Act
+    opc.manage_output_tags = AsyncMock(
+        return_value=(True, {'tag1': 0.1, 'tag2': 0.2}, False, None, False)
+    )
+
+    opc.process_confidence = MagicMock(return_value={'data': 'data'})
+    output_data, opc_metrics = await opc.write_opc_data(input_data)
+
+    # Assert
+    assert output_data == {'data': 'data'}
+    assert opc_metrics == {'server1': {'tag1': 0.1, 'tag2': 0.2}}
+    opc.manage_output_tags.assert_called_once_with(
+        'server1',
+        input_data['opc_output_config']['server1'],
+        mock_dataframe.return_value,
+        metadata['metadata'],
+    )
+    opc.process_confidence.assert_called_once_with(
+        mock_dataframe.return_value,
+        True,
+        metadata['metadata'],
+        session_bad=False,
+        opc_status=None,
+        reconnect_in_progress=False,
+    )
+
+
+@mark.asyncio
+async def test_write_opc_data_empty_config(opc):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
+        'opc_servers': ['server1'],
+        'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
+    }
+
+    # Act
+    await opc.write_opc_data(input_data)
+
+    # Assert
+    opc.opc_repository['server1'].write_data.assert_not_called()
+
+
+@mark.asyncio
+async def test_write_opc_data_no_validate_server(opc):
+    opc.validate_server = AsyncMock(return_value=False)
+    input_data = {
+        **metadata,
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
+        'opc_output_config': {
+            'server1': {
+                'prediction_tags': {'tag1': {'data_type': 'float'}},
+                'confidence_tags': {'tag2': {'data_type': 'float'}},
+            }
+        },
+    }
+
+    # Act
+    await opc.write_opc_data(input_data)
+
+    # Assert
+    opc.opc_repository['server1'].write_data.assert_not_called()
+
+
+@mark.parametrize(
+    'data,success,expected',
+    [
+        (DataFrame({'prediction_confidence': [0]}), True, 0),
+        (DataFrame({'prediction_confidence': [0]}), False, 12),
+    ],
+)
+def test_process_confidence(opc, data, success, expected):
+    result = opc.process_confidence(data, success, metadata['metadata'])
+    assert result['prediction_confidence'][0] == expected
+
+
+def test_process_confidence_session_bad(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    result = opc.process_confidence(
+        data,
+        False,
+        metadata['metadata'],
+        session_bad=True,
+        opc_status='BadSessionIdInvalid',
+    )
+    assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
+    assert result['comments'][0].startswith(OPC_SESSION_BAD_COMMENT_PREFIX)
+    assert 'BadSessionIdInvalid' in result['comments'][0]
+
+
+def test_process_confidence_generic_failure(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    result = opc.process_confidence(data, False, metadata['metadata'])
+    assert result['prediction_confidence'][0] == OPC_WRITTING_ERROR_CONFIDENCE
+    assert result['comments'][0] == OPC_WRITTING_ERROR_MESSAGE
+
+
+@mark.asyncio
+async def test_manage_output_tags_merges_error_flags(opc):
+    opc._write_tags_from_config = AsyncMock(
+        side_effect=[
+            ({'tag1': None}, True, 'BadSessionIdInvalid', False),
+            ({'tag2': 0.2}, False, None, True),
+        ]
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {
+        'prediction_tags': {'tag1': {'data_type': 'float'}},
+        'confidence_tags': {'tag2': {'data_type': 'float'}},
+    }
+
+    (
+        success,
+        metrics,
+        session_bad_seen,
+        opc_status,
+        reconnect_in_progress,
+    ) = await opc.manage_output_tags('server1', config, data, metadata['metadata'])
+
+    assert success is False
+    assert session_bad_seen is True
+    assert reconnect_in_progress is True
+    assert opc_status == 'BadSessionIdInvalid'
+    assert metrics == {'tag1': None, 'tag2': 0.2}
+
+
+def test_process_confidence_reconnect_in_progress(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    result = opc.process_confidence(
+        data,
+        False,
+        metadata['metadata'],
+        reconnect_in_progress=True,
+    )
+    assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
+    assert result['comments'][0] == OPC_RECONNECT_IN_PROGRESS_COMMENT
+
+
+def test_process_confidence_concatenates_multiple_comments(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    session_comment = f'{OPC_SESSION_BAD_COMMENT_PREFIX} BadSessionIdInvalid'
+
+    result = opc.process_confidence(
+        data,
+        False,
+        metadata['metadata'],
+        session_bad=True,
+        opc_status='BadSessionIdInvalid',
+        reconnect_in_progress=True,
+    )
+
+    assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
+    assert result['comments'][0] == OPC_COMMENT_SEPARATOR.join(
+        [session_comment, OPC_RECONNECT_IN_PROGRESS_COMMENT]
+    )
+
+
+@mark.asyncio
+async def test_validate_server(opc):
+    assert await opc.validate_server('server1', metadata) is True
+    assert await opc.validate_server('server2', metadata) is False
+
+
+@mark.asyncio
+async def test_close(opc):
+    opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
+    await opc.close()
+    opc.opc_repository['server1'].disconnect.assert_called_once()

tests/laborious/activities/test_storage.py

@@ -0,0 +1,323 @@
+import datetime
+import os
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
+from pytest import fixture, mark, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.activities.postgres import Postgres
+
+from laborious.activities.storage import Storage
+
+
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schedule_name': 'test_schedule',
+    }
+}
+
+
+@fixture
+@patch('laborious.activities.storage.MinioRepository')
+def storage(mock_minio_repository):
+    return Storage(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=mock_minio_repository.return_value,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___not_hasattr(mock_minio_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    minio_repo = mock_minio_repository.return_value
+    storage = Storage(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+    assert isinstance(storage, Postgres)
+
+    assert storage.minio_repository is minio_repo
+    mock_minio_repository.assert_not_called()
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___none_minio_repository(mock_minio_repository, storage):
+    storage.minio_repository = None
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    storage.__init__(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=None,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    assert storage.minio_repository is None
+    mock_minio_repository.assert_not_called()
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___done_repository(mock_minio_repository, storage):
+    storage.__init__(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=mock_minio_repository.return_value,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    mock_minio_repository.assert_not_called()
+    assert storage.minio_repository is not None
+
+
+def test_close(storage):
+    storage.minio_repository = MagicMock()
+
+    storage.close()
+
+    assert storage.minio_repository is None
+
+
+def test___del__(storage):
+    storage.close = MagicMock()
+
+    storage.__del__()
+
+    storage.close.assert_called_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_no_rows(storage):
+    storage.load_custom_query = AsyncMock(return_value=None)
+    storage_result = {'success': False}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=AsyncMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = await storage.load_query_with_minio_offload(
+            {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
+        )
+    assert result == storage_result
+    mock_from_dataframe.assert_awaited_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_inline(storage):
+    storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
+    storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=AsyncMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = await storage.load_query_with_minio_offload(
+            {
+                **metadata,
+                'query': 'SELECT 1',
+                'model_name': 'my-model',
+                'key_prefix': 'predictions/s',
+            }
+        )
+    assert result == storage_result
+    mock_from_dataframe.assert_awaited_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_minio(storage):
+    storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
+    storage_result = {'success': True, 'data': None, 'object_key': 'object-key'}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=AsyncMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = await storage.load_query_with_minio_offload(
+            {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
+        )
+
+    assert result == storage_result
+    mock_from_dataframe.assert_awaited_once()
+
+
+@mark.asyncio
+@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    storage.minio_repository.list_objects = AsyncMock(
+        return_value=[
+            'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
+            'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet',
+        ]
+    )
+    storage.minio_repository.delete_file = AsyncMock()
+    storage.send_notification_async = AsyncMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 1
+    assert result['failed_count'] == 0
+    deleted_key = (
+        'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
+    )
+    assert deleted_key in result['deleted']
+    assert result['deleted'][deleted_key]['success'] is True
+    storage.minio_repository.list_objects.assert_called_once_with(
+        prefix='training_datasets/m',
+        recursive=True,
+        metadata=metadata['metadata'],
+    )
+    storage.minio_repository.delete_file.assert_called_once_with(
+        object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
+        metadata=metadata['metadata'],
+    )
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_minio_not_initialized(storage):
+    storage.minio_repository = None
+
+    with raises(ValueError, match='Minio repository not initialized'):
+        await storage.load_query_with_minio_offload(
+            {**metadata, 'query': 'SELECT 1', 'model_name': 'm'}
+        )
+
+
+@mark.asyncio
+async def test_export_payload_to_postgres(storage):
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=MagicMock())
+    storage.export_data_to_postgres = AsyncMock(return_value={'success': True})
+
+    result = await storage.export_payload_to_postgres(
+        {**metadata, 'data': payload, 'schema': 'public', 'table': 't'}
+    )
+
+    payload.retrieve.assert_awaited_once_with(storage.minio_repository, metadata['metadata'])
+    storage.export_data_to_postgres.assert_awaited_once()
+    assert result == {'success': True}
+
+
+@mark.asyncio
+async def test_cleanup_minio_objects_expired_minio_not_initialized(storage):
+    storage.minio_repository = None
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'test'
+    with raises(ValueError, match='Minio repository not initialized'):
+        await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    storage.minio_repository.list_objects = AsyncMock(
+        return_value=['some/random/key-without-timestamp.parquet']
+    )
+    storage.minio_repository.delete_file = AsyncMock()
+    storage.send_notification_async = AsyncMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'test'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 0
+    storage.minio_repository.delete_file.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    old_key = 'training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
+    storage.minio_repository.list_objects = AsyncMock(return_value=[old_key])
+    storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error'))
+    storage.send_notification_async = AsyncMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 1
+    assert old_key in result['failed']
+    assert result['failed'][old_key]['success'] is False
+    assert result['failed'][old_key]['message'] == 'delete error'
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    storage.minio_repository.list_objects = AsyncMock(side_effect=Exception('list error'))
+    storage.send_notification_async = AsyncMock()
+    storage.error = MagicMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 0
+    storage.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
+        message='Error cleaning up MinIO objects: list error',
+        block='cleanup_minio_objects_expired',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+    storage.error.assert_called_once()