Code import - branch 0.5.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()

tests/laborious/utils/filters/test_conditional_filters.py

@@ -0,0 +1,44 @@
+from pandas import DataFrame
+
+from laborious.utils.filters.conditional_filters import (
+    filter_empty_data,
+    filter_specific_variables_null_values,
+)
+
+
+def test_filter_specific_variables_null_values():
+    assert (
+        filter_specific_variables_null_values(
+            DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
+            config={'variables': ['variable2']},
+        )
+        is False
+    )
+
+
+def test_filter_specific_variables_null_values_with_empty_data():
+    assert (
+        filter_specific_variables_null_values(DataFrame(), config={'variables': ['variable2']})
+        is False
+    )
+
+
+def test_filter_specific_variables_null_values_with_null_values():
+    assert (
+        filter_specific_variables_null_values(
+            DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
+            config={'variables': ['variable2']},
+        )
+        is True
+    )
+
+
+def test_filter_empty_data():
+    assert filter_empty_data(DataFrame(), {}) is True
+
+
+def test_filter_empty_data_with_data():
+    assert (
+        filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
+        is False
+    )

tests/laborious/utils/filters/test_mlflow_filters.py

@@ -0,0 +1,23 @@
+from pandas import DataFrame
+
+from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
+
+
+def test_api_error_filter_invalid_response():
+    assert api_error_filter(None, {}) is True  # NOSONAR
+
+
+def test_api_error_filter_valid_response_fail():
+    assert api_error_filter({'success': False}, {}) is True
+
+
+def test_api_error_filter_valid_response_success():
+    assert api_error_filter({'success': True}, {}) is False
+
+
+def test_nan_values_filter_all_nan_values():
+    assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) is True
+
+
+def test_nan_values_filter_no_nan_values():
+    assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False

tests/laborious/utils/models/test_minio_dataframe_payload.py

@@ -0,0 +1,266 @@
+from datetime import datetime
+from io import BytesIO
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from pandas import DataFrame
+
+from laborious.utils.models.minio_dataframe_payload import (
+    MinioDataFramePayload,
+    _build_object_key,
+)
+
+
+def test_parse_object_timestamp_hyphenated_model():
+    key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
+    ts = MinioDataFramePayload.parse_object_timestamp(key)
+    assert ts == datetime(2024, 6, 15, 10, 30, 45)
+
+
+def test_parse_object_timestamp_transform():
+    key = 'p/m-transform-2024-01-02_03-04-05.parquet'
+    ts = MinioDataFramePayload.parse_object_timestamp(key)
+    assert ts == datetime(2024, 1, 2, 3, 4, 5)
+
+
+def test_parse_object_timestamp_invalid():
+    assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
+
+
+def test_estimate_size_bytes_returns_positive_for_nonempty_frame():
+    df = DataFrame({'a': [1, 2]})
+    size = MinioDataFramePayload.estimate_size_bytes(df)
+    assert isinstance(size, int)
+    assert size > 0
+
+
+def test_cleanup_prefix_when_offloaded_returns_object_prefix():
+    payload = MinioDataFramePayload(
+        last_timestamp='t',
+        data=None,
+        object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
+        object_prefix='training_datasets/m',
+    )
+    assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m'
+
+
+def test_cleanup_prefix_when_inline_returns_none():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
+    assert MinioDataFramePayload.cleanup_prefix(payload) is None
+
+
+def test_has_data_true_when_object_key_set():
+    payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
+    assert payload.has_data() is True
+
+
+@pytest.mark.asyncio
+async def test_retrieve_inline_dict_as_dataframe():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
+    minio = AsyncMock()
+    out = await payload.retrieve(minio, {'metadata': {}})
+    assert list(out.columns) == ['a']
+    minio.download_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_retrieve_downloads_parquet_when_offloaded():
+    source = DataFrame({'a': [1, 2]})
+    buf = BytesIO()
+    source.to_parquet(buf, engine='pyarrow', index=True)
+    file_bytes = buf.getvalue()
+
+    payload = MinioDataFramePayload(
+        last_timestamp='t',
+        data=None,
+        object_key='training_datasets/m/f.parquet',
+        object_prefix='training_datasets/m',
+    )
+    minio = AsyncMock()
+    minio.download_file = AsyncMock(return_value=file_bytes)
+
+    out = await payload.retrieve(minio, {'metadata': {}})
+
+    minio.download_file.assert_awaited_once_with(
+        object_name='training_datasets/m/f.parquet',
+        metadata={'metadata': {}},
+    )
+    assert list(out.columns) == ['a']
+
+
+def test_build_object_key():
+    key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
+    assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
+    assert prefix == 'prediction_datasets/my-model'
+
+
+def test_build_object_key_strips_slashes():
+    key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
+    assert prefix == 'prediction_datasets/my-model'
+    assert key.startswith('prediction_datasets/my-model/')
+
+
+def test_estimate_size_bytes_fallback():
+    df = DataFrame({'a': [1, 2]})
+    with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')):
+        size = MinioDataFramePayload.estimate_size_bytes(df)
+    assert isinstance(size, int)
+    assert size > 0
+
+
+def test_parse_object_timestamp_bad_datetime():
+    key = 'p/m-initial-9999-99-99_99-99-99.parquet'
+    assert MinioDataFramePayload.parse_object_timestamp(key) is None
+
+
+@pytest.mark.asyncio
+async def test_retrieve_empty_when_no_data():
+    payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
+    minio = AsyncMock()
+    out = await payload.retrieve(minio, {})
+    assert out.empty
+    minio.download_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+async def test_from_dataframe_none(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = AsyncMock()
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=None,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+        status={'success': False, 'message': 'no data'},
+    )
+    assert result.data is None
+    assert result.status == {'success': False, 'message': 'no data'}
+    assert result.object_key is None
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+async def test_from_dataframe_empty(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = AsyncMock()
+    mock_df = MagicMock()
+    mock_df.__bool__ = MagicMock(return_value=True)
+    mock_df.empty = True
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=mock_df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+    )
+    assert result.data is None
+    assert result.object_key is None
+
+
+def _mock_dataframe(data_dict, timestamp_values=None):
+    """Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
+    mock_df = MagicMock()
+    mock_df.__bool__ = MagicMock(return_value=True)
+    mock_df.empty = False
+    if timestamp_values is None:
+        timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
+    ts_col = MagicMock()
+    ts_col.values.tolist.return_value = timestamp_values
+    mock_df.__getitem__ = MagicMock(return_value=ts_col)
+    mock_df.to_dict.return_value = data_dict
+    buf = BytesIO()
+    DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
+    mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
+    return mock_df
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
+async def test_from_dataframe_inline():
+    minio = AsyncMock()
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+    )
+    assert result.data is not None
+    assert result.object_key is None
+    assert result.last_timestamp == '2024-01-01'
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
+async def test_from_dataframe_offloaded(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = AsyncMock()
+    minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'})
+    minio.bucket = 'test-bucket'
+
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+        workflow_metadata={'wf': 'data'},
+    )
+    assert result.data is None
+    assert result.object_key == 'full/key.parquet'
+    assert result.bucket == 'test-bucket'
+    assert result.uri == 's3://test-bucket/full/key.parquet'
+    minio.upload_file.assert_awaited_once()
+
+
+def test_from_dict_inline():
+    raw = {
+        'last_timestamp': '2024-01-01T00:00:00+00:00',
+        'status': None,
+        'data': {'col1': {0: 'val1'}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
+    assert payload.data == {'col1': {0: 'val1'}}
+    assert payload.object_key is None
+
+
+def test_from_dict_offloaded():
+    raw = {
+        'last_timestamp': '2024-06-15T10:30:45+00:00',
+        'status': {'success': True},
+        'data': None,
+        'bucket': 'my-bucket',
+        'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+        'object_prefix': 'training_datasets/model',
+        'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.data is None
+    assert payload.bucket == 'my-bucket'
+    assert payload.object_key == raw['object_key']
+    assert payload.object_prefix == 'training_datasets/model'
+    assert payload.uri == raw['uri']
+    assert payload.status == {'success': True}
+
+
+def test_from_dict_minimal_keys():
+    raw = {'last_timestamp': '2024-01-01'}
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert payload.last_timestamp == '2024-01-01'
+    assert payload.data is None
+    assert payload.bucket is None
+    assert payload.object_key is None
+
+
+def test_from_dict_passthrough_existing_instance():
+    original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
+    result = MinioDataFramePayload.from_dict(original)
+    assert result is original

tests/laborious/utils/repository/test_opc_repository.py

@@ -0,0 +1,610 @@
+import asyncio
+import json
+from datetime import datetime
+from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
+
+import pytest
+from asyncua.crypto.security_policies import SecurityPolicyBasic256
+from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.utils.repository.opc_repository import (
+    OpcClientAlreadyExistsError,
+    OpcClientNotInitializedError,
+    OpcRepository,
+    OpcSessionAlreadyConnectedError,
+    is_reconnectable_opcua_bad,
+)
+
+
+@pytest.fixture
+def mock_logger():
+    return Mock()
+
+
+@pytest.fixture
+def opc_repository(mock_logger):
+    repository = OpcRepository(
+        opc_id='test_repo',
+        server_name='test_server',
+        url='opc.tcp://localhost:4840',
+        logger=mock_logger,
+        notification_handler=Mock(),
+        reconnection_interval=60,
+        server_uri='urn:test:server',
+        cert_path='/path/to/cert.pem',
+        private_key_path='/path/to/key.pem',
+        server_cert_path='/path/to/server_cert.pem',
+        metrics_controller=AsyncMock(),
+    )
+    repository.disconnection_interval = 0.1
+    repository.send_notification = MagicMock()
+    repository.send_notification_async = AsyncMock()
+    repository.emit_metric = AsyncMock()
+    repository.info = MagicMock()
+    repository.error = MagicMock()
+    repository.warning = MagicMock()
+    repository.debug = MagicMock()
+    repository._session_ready.set()
+    return repository
+
+
+@pytest.fixture
+def mock_client():
+    with patch('laborious.utils.repository.opc_repository.Client') as mock:
+        client_instance = AsyncMock()
+        mock.return_value = client_instance
+        yield client_instance
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+def test_init(opc_repository):
+    assert opc_repository.id == 'test_repo'
+    assert opc_repository.server_name == 'test_server'
+    assert opc_repository.url == 'opc.tcp://localhost:4840'
+    assert opc_repository.server_uri == 'urn:test:server'
+    assert opc_repository.cert_path == '/path/to/cert.pem'
+    assert opc_repository.private_key_path == '/path/to/key.pem'
+    assert opc_repository.server_cert_path == '/path/to/server_cert.pem'
+    assert opc_repository.reconnection_interval == 60
+    assert opc_repository.client is None
+    assert opc_repository.last_reconnection_time is None
+
+
+@pytest.mark.asyncio
+async def test_set_security(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    await opc_repository.set_security()
+
+    mock_client.application_uri = 'urn:test:server'
+    mock_client.set_security.assert_called_once_with(
+        SecurityPolicyBasic256,
+        certificate='/path/to/cert.pem',
+        private_key='/path/to/key.pem',
+        server_certificate='/path/to/server_cert.pem',
+    )
+    assert mock_client.secure_channel_timeout == 600_000
+    assert mock_client.session_timeout == 600_000
+
+
+@pytest.mark.asyncio
+async def test_set_security_missing_certificates(opc_repository):
+    opc_repository.cert_path = None
+    opc_repository.private_key_path = None
+
+    try:
+        await opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
+
+
+@pytest.mark.asyncio
+async def test_set_security_missing_client(opc_repository):
+    opc_repository.client = None
+    try:
+        await opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Client must be initialized before setting security'
+
+
+@pytest.mark.asyncio
+async def test_connect_with_security(opc_repository, mock_client):
+    opc_repository._create_client = AsyncMock()
+    opc_repository._open_session = AsyncMock(return_value=(True, {}))
+    result = await opc_repository.connect()
+
+    opc_repository._create_client.assert_called_once()
+    opc_repository._open_session.assert_called_once()
+    assert result == (True, {})
+
+
+@pytest.mark.asyncio
+async def test_connect_without_security(opc_repository, mock_client):
+    opc_repository.cert_path = None
+    opc_repository._create_client = AsyncMock()
+    opc_repository._open_session = AsyncMock(return_value=(True, {}))
+    opc_repository.set_security = AsyncMock()
+    result = await opc_repository.connect()
+
+    opc_repository._create_client.assert_called_once()
+    opc_repository._open_session.assert_called_once()
+    opc_repository.set_security.assert_not_called()
+    assert result == (True, {})
+
+
+@pytest.mark.asyncio
+async def test_connect_raises_when_session_already_open(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'open'
+    mock_client.uaclient = MagicMock(protocol=proto)
+
+    with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
+        await opc_repository.connect()
+
+
+@pytest.mark.asyncio
+async def test_create_client_raises_when_client_exists(opc_repository, mock_client):
+    opc_repository.client = mock_client
+
+    with pytest.raises(OpcClientAlreadyExistsError, match='already exists'):
+        await opc_repository._create_client()
+
+
+@pytest.mark.asyncio
+async def test_open_session_success(opc_repository):
+    closed_proto = MagicMock()
+    closed_proto.state = 'closed'
+    opc_repository.client = AsyncMock()
+    opc_repository.client.uaclient = MagicMock(protocol=closed_proto)
+    opc_repository.client.session_timeout = 600_000
+    opc_repository.client.secure_channel_timeout = 600_000
+
+    open_proto = MagicMock()
+    open_proto.state = 'open'
+    open_proto.authentication_token = 'tok'
+
+    async def connect_side_effect():
+        opc_repository.client.uaclient.protocol = open_proto
+
+    opc_repository.client.connect = AsyncMock(side_effect=connect_side_effect)
+
+    result = await opc_repository._open_session()
+
+    opc_repository.client.connect.assert_called_once()
+    assert opc_repository.last_reconnection_time is None
+    assert result == (True, {})
+    assert opc_repository._session_ready.is_set()
+
+
+@pytest.mark.asyncio
+async def test_open_session_raises_when_already_connected(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'open'
+    mock_client.uaclient = MagicMock(protocol=proto)
+
+    with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
+        await opc_repository._open_session()
+
+
+@pytest.mark.asyncio
+async def test_open_session_fail(opc_repository):
+    opc_repository._disconnect_locked = AsyncMock()
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient = MagicMock(protocol=MagicMock(state='closed'))
+    opc_repository.client.connect = AsyncMock(side_effect=Exception('Test error'))
+
+    is_connected, error_data = await opc_repository._open_session()
+
+    opc_repository._disconnect_locked.assert_called_once()
+    opc_repository.client.connect.assert_called_once()
+    assert is_connected is False
+    assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'
+    assert error_data['message'] == 'Failed to connect to OPC server: Test error'
+    assert error_data['block'] == 'opc_repository'
+    assert error_data['level'] == NotificationLevel.ERROR
+    assert error_data['attachment_content'] is not None
+
+
+@pytest.mark.asyncio
+async def test_open_session_raises_when_no_client(opc_repository):
+    opc_repository.client = None
+
+    with pytest.raises(OpcClientNotInitializedError, match='not initialized'):
+        await opc_repository._open_session()
+
+
+@pytest.mark.asyncio
+async def test_disconnection_fallback_success(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    mock_client.disconnect.return_value = True
+    result = await opc_repository._disconnection_fallback()
+
+    mock_client.disconnect.assert_called_once()
+    assert result == []
+
+
+@pytest.mark.asyncio
+async def test_disconnection_fallback_fail(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    mock_client.disconnect.side_effect = Exception('Test error')
+    result = await opc_repository._disconnection_fallback()
+    assert result == [
+        {'attempt': 1, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 2, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 3, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 4, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 5, 'error': 'Test error', 'traceback': ANY},
+    ]
+    assert mock_client.disconnect.call_count == 5
+
+
+@pytest.mark.asyncio
+async def test_disconnect(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    opc_repository._disconnection_fallback = AsyncMock(return_value=[])
+    await opc_repository.disconnect()
+
+    opc_repository._disconnection_fallback.assert_called_once()
+    assert opc_repository.client is None
+    assert opc_repository._allow_reconnect is False
+
+
+@pytest.mark.asyncio
+async def test_disconnect_no_client(opc_repository):
+    opc_repository.client = None
+    assert await opc_repository.disconnect() is None
+
+
+@pytest.mark.asyncio
+async def test_disconnect_error(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    opc_repository._disconnection_fallback = AsyncMock(
+        return_value=[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}]
+    )
+    await opc_repository.disconnect()
+
+    opc_repository._disconnection_fallback.assert_called_once()
+    opc_repository.send_notification_async.assert_called_once_with(
+        metadata=opc_repository.metadata,
+        notification_id=f'OPC_DISCONNECTION_ERROR_{opc_repository.id}',
+        message='Failed to disconnect from OPC server in 5 attempts.',
+        block='opc_repository',
+        level=NotificationLevel.ERROR,
+        attachment_content=json.dumps(
+            [{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}], indent=4
+        ),
+    )
+    assert opc_repository.client is None
+
+
+@pytest.mark.asyncio
+async def test_validate_connection_none_client(opc_repository):
+    opc_repository.client = None
+    response = await opc_repository.validate_connection()
+    assert response == (False, opc_repository._not_connected_error())
+
+
+@pytest.mark.asyncio
+async def test_validate_connection_session_not_open(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = None
+
+    response = await opc_repository.validate_connection()
+
+    assert response == (False, opc_repository._not_connected_error())
+    opc_repository.error.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_validate_connection_success(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock()
+    opc_repository.client.uaclient.protocol.state = 'open'
+
+    output = await opc_repository.validate_connection()
+    assert output == (True, {})
+
+
+@pytest.mark.asyncio
+async def test_write_data_validate_connection_do_nothing(opc_repository):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = AsyncMock(get_node=MagicMock())
+    mock_node = AsyncMock()
+    opc_repository.client.get_node.return_value = mock_node
+
+    result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+
+    opc_repository.validate_connection.assert_called_once()
+    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
+    assert result == (True, {'response_time': ANY})
+
+
+@pytest.mark.asyncio
+async def test_write_data_validate_connection_failed(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock(state='closed')
+    opc_repository._start_reconnect = AsyncMock()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository._start_reconnect.assert_called_once()
+    assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+    assert error_data['opc_status'] == 'ProtocolClosed'
+
+
+@pytest.mark.asyncio
+async def test_write_data_get_node_failed(opc_repository):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = AsyncMock()
+    opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository.validate_connection.assert_called_once()
+    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
+    assert is_success is False
+    assert error_data['notification_id'] == f'OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}'
+    assert (
+        error_data['message']
+        == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
+    assert error_data['level'] == NotificationLevel.ERROR
+    assert error_data['attachment_content'] is not None
+
+
+@pytest.mark.asyncio
+async def test_write_data_invalid_data_type(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'invalid_type', metadata['metadata']
+    )
+
+    opc_repository.validate_connection.assert_called_once()
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
+
+    assert is_success is False
+    assert error_data['notification_id'] == f'OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}'
+    assert (
+        error_data['message']
+        == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
+    assert error_data['level'] == NotificationLevel.ERROR
+    assert error_data.get('attachment_content') is None
+
+
+@pytest.mark.asyncio
+async def test_write_data(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+
+    result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
+    mock_node.write_value.assert_called_once()
+    assert result == (True, {'response_time': ANY})
+
+
+@pytest.mark.asyncio
+async def test_write_data_write_value_failed(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = Exception('Test error')
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository.validate_connection.assert_called_once()
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
+    mock_node.write_value.assert_called_once()
+    assert is_success is False
+    assert error_data['notification_id'] == f'OPC_WRITE_DATA_ERROR_{opc_repository.id}'
+    assert (
+        error_data['message']
+        == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
+    assert error_data['level'] == NotificationLevel.ERROR
+    assert error_data['attachment_content'] is not None
+
+
+def test_is_reconnectable_opcua_bad():
+    assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
+    assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
+    assert is_reconnectable_opcua_bad(Exception('other')) is False
+
+
+@pytest.mark.asyncio
+async def test_write_data_bad_session_id_invalid_schedules_reconnect(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    opc_repository._start_reconnect = AsyncMock()
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = BadSessionIdInvalid()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    mock_node.write_value.assert_called_once()
+    opc_repository._start_reconnect.assert_called_once()
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'session_bad'
+    assert error_data['opc_status'] == 'BadSessionIdInvalid'
+
+
+@pytest.mark.asyncio
+async def test_write_data_reconnect_in_progress_immediate(opc_repository):
+    opc_repository._session_ready.clear()
+    opc_repository._reconnect_task = asyncio.create_task(asyncio.sleep(60))
+    opc_repository.validate_connection = AsyncMock()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository._reconnect_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await opc_repository._reconnect_task
+    opc_repository._reconnect_task = None
+
+    opc_repository.validate_connection.assert_not_called()
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'reconnect_in_progress'
+
+
+@pytest.mark.asyncio
+async def test_start_reconnect_skips_within_interval(opc_repository):
+    opc_repository.last_reconnection_time = datetime.now()
+    opc_repository.reconnection_interval = 3600
+
+    await opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
+
+    assert opc_repository._reconnect_task is None
+
+
+@pytest.mark.asyncio
+async def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock(state='closed')
+    opc_repository._start_reconnect = AsyncMock()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository._start_reconnect.assert_called_once()
+    assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+    assert error_data['opc_status'] == 'ProtocolClosed'
+
+
+@pytest.mark.asyncio
+async def test_write_data_protocol_closed_skips_reconnect_within_interval(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock(state='closed')
+    opc_repository.last_reconnection_time = datetime.now()
+    opc_repository.reconnection_interval = 3600
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    assert opc_repository._reconnect_task is None
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+
+
+@pytest.mark.asyncio
+async def test_write_data_after_failed_reconnect_schedules_again(opc_repository):
+    opc_repository._session_ready.clear()
+    opc_repository.reconnection_interval = 0
+    opc_repository.last_reconnection_time = None
+    opc_repository._reconnect_locked = AsyncMock(
+        return_value=(False, {'message': 'connect failed'})
+    )
+
+    await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+    await asyncio.sleep(0.1)
+    assert opc_repository._reconnect_locked.call_count == 1
+    assert not opc_repository._reconnect_task_in_progress()
+
+    await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+    await asyncio.sleep(0.1)
+    assert opc_repository._reconnect_locked.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_write_data_after_disconnect_does_not_schedule_reconnect(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'closed'
+    mock_client.uaclient = MagicMock(protocol=proto)
+    opc_repository._disconnection_fallback = AsyncMock(return_value=[])
+    await opc_repository.disconnect()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    assert opc_repository._reconnect_task is None
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+
+
+@pytest.mark.asyncio
+async def test_parallel_bad_writes_single_reconnect_task(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    opc_repository.reconnection_interval = 0
+    opc_repository.last_reconnection_time = None
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = BadSessionIdInvalid()
+
+    connect_count = 0
+
+    async def slow_reconnect():
+        nonlocal connect_count
+        connect_count += 1
+        await asyncio.sleep(0.05)
+        opc_repository._session_ready.set()
+        return True, {}
+
+    opc_repository._reconnect_locked = slow_reconnect
+
+    results = await asyncio.gather(
+        opc_repository.write_data('ns=2;s=TestNode', 1.0, 'float', metadata['metadata']),
+        opc_repository.write_data('ns=2;s=TestNode2', 2.0, 'float', metadata['metadata']),
+    )
+    await asyncio.sleep(0.15)
+
+    assert connect_count <= 1
+    assert 1 <= mock_node.write_value.call_count <= 2
+    error_kinds = [r[1].get('opc_error_kind') for r in results]
+    assert error_kinds.count('session_bad') >= 1
+    assert all(k in ('session_bad', 'reconnect_in_progress') for k in error_kinds)
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.repository.opc_repository.datetime')
+async def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_repository):
+    mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 12, 0, 0))
+    opc_repository._disconnect_locked = AsyncMock()
+    opc_repository._connect_locked = AsyncMock(return_value=(True, {}))
+
+    result = await opc_repository._reconnect_locked()
+
+    opc_repository._disconnect_locked.assert_called_once()
+    opc_repository._connect_locked.assert_called_once()
+    assert result == (True, {})
+    assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)

tests/laborious/utils/test_connectors_config.py

@@ -0,0 +1,122 @@
+from os import environ
+
+from laborious.utils.connectors_config import (
+    build_minio_config,
+    build_mlflow_config,
+    build_opc_config,
+)
+
+
+def test_build_mlflow_config_with_env_vars():
+    # Arrange
+    environ['MLFLOW_HOST'] = 'http://test-host'
+    environ['MLFLOW_PORT'] = '8080'
+    environ['MLFLOW_USERNAME'] = 'test-user'
+    environ['MLFLOW_PASSWORD'] = 'test-pass'
+
+    # Act
+    config = build_mlflow_config()
+
+    # Assert
+    assert config['host'] == 'http://test-host'
+    assert config['port'] == 8080
+    assert config['username'] == 'test-user'
+    assert config['password'] == 'test-pass'
+
+
+def test_build_mlflow_config_with_defaults():
+    # Arrange
+    # Clear any existing env vars
+    environ.pop('MLFLOW_HOST', None)
+    environ.pop('MLFLOW_PORT', None)
+    environ.pop('MLFLOW_USERNAME', None)
+    environ.pop('MLFLOW_PASSWORD', None)
+
+    # Act
+    config = build_mlflow_config()
+
+    # Assert
+    assert config['host'] == 'http://localhost'
+    assert config['port'] == 5080
+    assert config['username'] == 'aignosi'
+    assert config['password'] == 'aignosi'
+
+
+def test_build_opc_config_with_env_vars():
+    # Arrange
+    environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
+
+    # Act
+    config = build_opc_config()
+
+    # Assert
+    assert config['opc']['name'] == 'test-opc'
+    assert config['opc']['url'] == 'opc.tcp://test:4840'
+
+
+def test_build_opc_config_with_individual_env_vars():
+    # Arrange
+    environ.pop('OPC_CONFIG', None)
+    environ['OPC_ID'] = '1'
+    environ['OPC_URL'] = 'opc.tcp://test:4840'
+    environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
+    environ['OPC_RECONNECTION_INTERVAL'] = '300'
+
+    # Act
+    config = build_opc_config()
+
+    # Assert
+    assert config['1']['id'] == '1'
+    assert config['1']['url'] == 'opc.tcp://test:4840'
+    assert config['1']['server_uri'] == 'opc.tcp://test:4840'
+    assert config['1']['reconnection_interval'] == 300
+
+
+def test_build_opc_config_with_defaults():
+    # Arrange
+    environ.pop('OPC_CONFIG', None)
+    environ.pop('OPC_ID', None)
+    environ.pop('OPC_URL', None)
+    environ.pop('OPC_SERVER_URI', None)
+    environ.pop('OPC_RECONNECTION_INTERVAL', None)
+
+    # Act
+    config = build_opc_config()
+
+    # Assert
+    assert config['1']['id'] == '1'
+    assert config['1']['url'] == 'opc.tcp://localhost:4840'
+    assert config['1']['server_uri'] == 'opc.tcp://localhost:4840'
+    assert config['1']['reconnection_interval'] == 120
+
+
+def test_build_minio_config_with_env_vars():
+    environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
+    environ['MINIO_ACCESS_KEY'] = 'test-key'
+    environ['MINIO_SECRET_KEY'] = 'test-secret'
+    environ['MINIO_REGION_NAME'] = 'test-region'
+    environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
+    assert build_minio_config() == {
+        'endpoint_url': 'http://test-host',
+        'access_key': 'test-key',
+        'secret_key': 'test-secret',
+        'default_bucket': 'test-bucket',
+        'retention_hours': 24,
+        'secure': False,
+    }
+
+
+def test_build_minio_config_with_defaults():
+    environ.pop('MINIO_ENDPOINT_URL', None)
+    environ.pop('MINIO_ACCESS_KEY', None)
+    environ.pop('MINIO_SECRET_KEY', None)
+    environ.pop('MINIO_REGION_NAME', None)
+    environ.pop('MINIO_DEFAULT_BUCKET', None)
+    assert build_minio_config() == {
+        'endpoint_url': 'http://localhost:9000',
+        'access_key': 'minioadmin',
+        'secret_key': 'minioadmin',
+        'default_bucket': 'laborious',
+        'retention_hours': 24,
+        'secure': False,
+    }

tests/laborious/worker/test_runtime_task_queues.py

@@ -0,0 +1,14 @@
+from sientia_do.temporal.worker.prepare_worker import build_queue_name
+
+from laborious.workflows.drift import Drift
+from laborious.workflows.minimal_retrain import MinimalRetrain
+from laborious.workflows.predictions_batch import PredictionsBatch
+from laborious.workflows.simple_metrics import SimpleMetrics
+
+
+def test_runtime_scoped_queue_names():
+    runtime = 'prod-a'
+    assert build_queue_name(PredictionsBatch.__name__, runtime) == 'predictions_batch-prod-a-queue'
+    assert build_queue_name(MinimalRetrain.__name__, runtime) == 'minimal_retrain-prod-a-queue'
+    assert build_queue_name(Drift.__name__, runtime) == 'drift-prod-a-queue'
+    assert build_queue_name(SimpleMetrics.__name__, runtime) == 'simple_metrics-prod-a-queue'

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -0,0 +1,679 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.activities import Activities
+from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
+
+
+@fixture
+def format_and_export_prediction():
+    return FormatAndExportPrediction()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'opc_servers': ['test_server'],
+        'opc_output_config': {'test': 'config'},
+        'prediction_store_policy': 'erl:1',
+    }
+
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1
+
+
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_none_path_flag_with_transformed_data(
+    workflow_mock, format_and_export_prediction
+):
+    # Arrange
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'transformed_data': {'transformed': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0.9,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'opc_servers': ['test_server'],
+        'opc_output_config': {'test': 'config'},
+        'prediction_store_policy': 'lts:1',
+    }
+
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+    transformed_data = MagicMock()
+
+    workflow_mock.execute_local_activity_method.side_effect = [
+        prediction_data,  # format_prediction
+        transformed_data,  # format_transformed_data
+    ]
+
+    write_transformed_handler = AsyncMock()
+    workflow_mock.start_activity_method.return_value = write_transformed_handler
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),  # write_opc_data
+        MagicMock(),  # export_data_to_postgres (prediction)
+        MagicMock(),  # write_metrics
+    ]
+
+    # Act
+    await format_and_export_prediction.run(input_data)
+
+    # Assert - format_prediction call
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.format_transformed_data,
+                {
+                    **metadata,
+                    'data': input_data['transformed_data'],
+                    'model_id': input_data['model_id'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    # Assert - start_activity_method for transformed data export
+    workflow_mock.start_activity_method.assert_called_once_with(
+        Activities.export_payload_to_postgres,
+        {
+            **metadata,
+            'schema': input_data['schema'],
+            'table_name': input_data['transform_table_name'],
+            'data': transformed_data,
+            'timestamp_conversion': {
+                'column': 'timestamp',
+                'format': DATETIME_FORMAT_WITH_TZ,
+            },
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    # Assert - write_opc_data call
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': prediction_data,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    # Assert - export_data_to_postgres for prediction call
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    # Assert - write_metrics call
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    # Assert - verify counts
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 2
+    assert workflow_mock.start_activity_method.call_count == 1
+
+
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': 'default',
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'opc_servers': ['test_server'],
+        'opc_output_config': {'test': 'config'},
+        'comment': 'test_comment',
+    }
+
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_default_prediction,
+                {
+                    **metadata,
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'comment': input_data['comment'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1
+
+
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+        'prediction_store_policy': 'erl:1',
+    }
+
+    pi_web_api_data = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        pi_web_api_data,  # write_pi_web_api_data
+        MagicMock(),  # export_data_to_postgres
+        MagicMock(),  # write_metrics
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': pi_web_api_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': pi_web_api_data,
+                    'opc_metrics': {},
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1
+
+
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_none_path_flag_with_pi_web_api_and_opc(
+    workflow_mock, format_and_export_prediction
+):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'opc_output_config': {'test': 'config'},
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+        'prediction_store_policy': 'erl:1',
+    }
+
+    prediction_data = MagicMock()
+    pi_web_api_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        pi_web_api_data,  # write_pi_web_api_data
+        (prediction_data, opc_metrics),  # write_opc_data
+        MagicMock(),  # export_data_to_postgres
+        MagicMock(),  # write_metrics
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': pi_web_api_data,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 4
+    assert workflow_mock.execute_local_activity_method.call_count == 1
+
+
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': 'default',
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+        'comment': 'test_comment',
+    }
+
+    pi_web_api_data = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        pi_web_api_data,  # write_pi_web_api_data
+        MagicMock(),  # export_data_to_postgres
+        MagicMock(),  # write_metrics
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': pi_web_api_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -0,0 +1,842 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+from pytest import fixture, mark
+
+from laborious.activities.activities import Activities
+from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
+
+
+@fixture
+def prediction_process():
+    return PredictionProcess()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+        'schedule_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock, prediction_process):
+    prediction_process.path_flag_handler = AsyncMock(return_value=False)
+    # Arrange
+    data_payload = MagicMock()
+    data_payload.cleanup_prefix.return_value = 'training_datasets/test'
+    data_payload.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    input_data = {
+        'metadata': metadata,
+        'data': data_payload,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'model_id': 1,
+        'input_filters': {'test': 'filter'},
+        'mlflow_transform_filters': {'test': 'filter'},
+        'mlflow_predict_filters': {'test': 'filter'},
+        'model_name': 'test_model_name',
+        'model_config': {'retention': '30'},
+        'path_priority': ['continue', 'repeat', 'stop'],
+        'opc_output_config': {'test': 'config'},
+        'pi_web_api_output_config': {'test': 'config'},
+        'prediction_store_policy': 'lts:1',
+    }
+
+    # Mock the activity responses
+    workflow_mock.execute_activity_method.side_effect = [
+        {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+        {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+        MagicMock(),
+    ]
+    workflow_mock.execute_local_activity_method.side_effect = [
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
+        # mlflow_response_gate (transform)
+        ('continue', 0.95, 'Error'),
+        # mlflow_content_gate (transform)
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
+        # mlflow_response_gate (predict)
+        ('continue', 0.95, 'Error'),
+    ]
+
+    # Act
+    await prediction_process.run(input_data)
+
+    # Assert
+    assert workflow_mock.execute_local_activity_method.call_count == 4
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    **metadata,
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    **metadata,
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
+                    **metadata,
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_predict,
+                {
+                    **metadata,
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    **metadata,
+                    'filters': input_data['mlflow_predict_filters'],
+                    'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+                    'type': 'predict',
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_child_workflow.assert_called_once_with(
+        'subworkflow.format_and_export_prediction',
+        {
+            'metadata': metadata,
+            'on_conflict': 'error',
+            'path_flag': 'continue',
+            'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+            'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+            'prediction_confidence': 0.95,
+            'timestamp': '2024-01-01',
+            'model_id': 1,
+            'model_name': 'test_model_name',
+            'model_config': input_data['model_config'],
+            'opc_output_config': input_data['opc_output_config'],
+            'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+            'schema': input_data['schema'],
+            'table_name': input_data['table_name'],
+            'transform_table_name': input_data['transform_table_name'],
+            'comment': 'Error',
+            'prediction_store_policy': input_data['prediction_store_policy'],
+        },
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
+    prediction_process.path_flag_handler = AsyncMock(return_value=True)
+    # Arrange
+    data_payload = MagicMock()
+    data_payload.cleanup_prefix.return_value = 'training_datasets/test'
+    data_payload.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    input_data = {
+        'metadata': metadata,
+        'data': data_payload,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'model_id': 1,
+        'input_filters': {'test': 'filter'},
+        'mlflow_transform_filters': {'test': 'filter'},
+        'mlflow_predict_filters': {'test': 'filter'},
+        'model_name': 'test_model_name',
+        'model_config': {'retention': '30'},
+        'path_priority': ['continue', 'repeat', 'stop'],
+        'opc_output_config': {'test': 'config'},
+    }
+
+    # Mock the activity responses
+    workflow_mock.execute_local_activity_method.side_effect = [
+        ('stop', 0.95, 'Input data with bad quality'),  # input_gate
+    ]
+
+    # Act
+    await prediction_process.run(input_data)
+
+    # Assert
+    assert workflow_mock.execute_local_activity_method.call_count == 1
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
+    # Arrange
+    data_payload = MagicMock()
+    data_payload.cleanup_prefix.return_value = 'training_datasets/test'
+    data_payload.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    input_data = {
+        'metadata': metadata,
+        'data': data_payload,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'model_id': 1,
+        'input_filters': {'test': 'filter'},
+        'mlflow_transform_filters': {'test': 'filter'},
+        'mlflow_predict_filters': {'test': 'filter'},
+        'model_name': 'test_model_name',
+        'model_config': {'retention': '30'},
+        'path_priority': ['continue', 'repeat', 'stop'],
+        'opc_output_config': {'test': 'config'},
+    }
+
+    # Mock the activity responses
+    workflow_mock.execute_activity_method.side_effect = [
+        {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+        MagicMock(),
+    ]
+    workflow_mock.execute_local_activity_method.side_effect = [
+        ('repeat', 0.95, 'Input data with bad quality'),  # input_gate
+        ('continue', 0.95, 'Error'),  # mlflow_response_gate (transform)
+    ]
+
+    # Act
+    await prediction_process.run(input_data)
+
+    # Assert
+    assert workflow_mock.execute_local_activity_method.call_count == 2
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
+    # Arrange
+    data_payload = MagicMock()
+    data_payload.cleanup_prefix.return_value = 'training_datasets/test'
+    data_payload.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    input_data = {
+        'metadata': metadata,
+        'data': data_payload,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'model_id': 1,
+        'input_filters': {'test': 'filter'},
+        'mlflow_transform_filters': {'test': 'filter'},
+        'mlflow_predict_filters': {'test': 'filter'},
+        'model_name': 'test_model_name',
+        'model_config': {'retention': '30'},
+        'path_priority': ['continue', 'repeat', 'stop'],
+        'opc_output_config': {'test': 'config'},
+    }
+
+    # Mock the activity responses
+    workflow_mock.execute_activity_method.side_effect = [
+        {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+        MagicMock(),
+    ]
+    workflow_mock.execute_local_activity_method.side_effect = [
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
+        # mlflow_response_gate (transform)
+        ('continue', 0.95, 'Error'),
+        # mlflow_content_gate (transform)
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
+    ]
+
+    # Act
+    await prediction_process.run(input_data)
+
+    # Assert
+    assert workflow_mock.execute_local_activity_method.call_count == 3
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
+    # Arrange
+    data_payload = MagicMock()
+    data_payload.cleanup_prefix.return_value = 'training_datasets/test'
+    data_payload.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    input_data = {
+        'metadata': metadata,
+        'data': data_payload,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'model_id': 1,
+        'input_filters': {'test': 'filter'},
+        'mlflow_transform_filters': {'test': 'filter'},
+        'mlflow_predict_filters': {'test': 'filter'},
+        'model_name': 'test_model_name',
+        'model_config': {'retention': '30'},
+        'path_priority': ['continue', 'repeat', 'stop'],
+        'opc_output_config': {'test': 'config'},
+    }
+
+    # Mock the activity responses
+    workflow_mock.execute_activity_method.side_effect = [
+        {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+        {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+        MagicMock(),
+    ]
+    workflow_mock.execute_local_activity_method.side_effect = [
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
+        # mlflow_response_gate (transform)
+        ('continue', 0.95, 'Error'),
+        # mlflow_content_gate (transform)
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
+        ('continue', 0.95, 'Error'),  # mlflow_response_gate (predict)
+    ]
+
+    # Act
+    await prediction_process.run(input_data)
+
+    # Assert
+    assert workflow_mock.execute_local_activity_method.call_count == 4
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_predict,
+                {
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_predict_filters'],
+                    'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+                    'type': 'predict',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_path_flag_handler_stop(workflow_mock, prediction_process):
+    # Arrange
+    data = {'test': 'data'}
+    path_flag = 'STOP'
+    confidence = 0.95
+    schema = 'test_schema'
+    table_name = 'test_table'
+    model = 'test_model'
+    last_timestamp = '2024-01-01'
+    model_name = 'test_model_name'
+    model_config = {'retention': '30'}
+
+    # Act
+    result = await prediction_process.path_flag_handler(
+        data,
+        path_flag,
+        {
+            'metadata': metadata,
+            'schema': schema,
+            'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_config': model_config,
+        },
+        confidence,
+        last_timestamp,
+        '',
+    )
+
+    # Assert
+    assert result is True
+    workflow_mock.execute_local_activity_method.assert_not_called()
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
+    # Arrange
+    data = {'test': 'data'}
+    path_flag = 'repeat'
+    confidence = 0.95
+    schema = 'test_schema'
+    table_name = 'test_table'
+    model = 'test_model'
+    last_timestamp = '2024-01-01'
+    model_name = 'test_model_name'
+    model_config = {'retention': '30'}
+
+    # Act
+    result = await prediction_process.path_flag_handler(
+        data,
+        path_flag,
+        {
+            'metadata': metadata,
+            'schema': schema,
+            'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_config': model_config,
+        },
+        confidence,
+        last_timestamp,
+        '',
+    )
+
+    # Assert
+    assert result is True
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.repeat_last_prediction,
+        {
+            **metadata,
+            'schema': schema,
+            'table_name': table_name,
+            'model': model,
+            'last_timestamp': last_timestamp,
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_path_flag_handler_continue(workflow_mock, prediction_process):
+    # Arrange
+    data = {'test': 'data'}
+    path_flag = 'CONTINUE'
+    confidence = 0.95
+    schema = 'test_schema'
+    table_name = 'test_table'
+    model = 'test_model'
+    last_timestamp = '2024-01-01'
+    model_name = 'test_model_name'
+    model_config = {'retention': '30'}
+    prediction_store_policy = 'erl:1'
+
+    # Act
+    result = await prediction_process.path_flag_handler(
+        data,
+        path_flag,
+        {
+            'metadata': metadata,
+            'schema': schema,
+            'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_config': model_config,
+            'opc_output_config': {'test': 'config'},
+            'pi_web_api_output_config': {
+                'endpoint': 'https://test-pi-server.com',
+                'prediction_tags': {},
+                'confidence_tags': {},
+            },
+            'prediction_store_policy': prediction_store_policy,
+        },
+        confidence,
+        last_timestamp,
+        'Prediction Process',
+    )
+
+    # Assert
+    assert result is True
+    workflow_mock.execute_activity_method.assert_not_called()
+    workflow_mock.execute_child_workflow.assert_called_once_with(
+        'subworkflow.format_and_export_prediction',
+        {
+            'metadata': metadata,
+            'path_flag': path_flag,
+            'data': data,
+            'prediction_confidence': confidence,
+            'timestamp': last_timestamp,
+            'model_id': model,
+            'model_name': model_name,
+            'model_config': model_config,
+            'schema': schema,
+            'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
+            'comment': 'Prediction Process',
+            'opc_output_config': {'test': 'config'},
+            'pi_web_api_output_config': {
+                'endpoint': 'https://test-pi-server.com',
+                'prediction_tags': {},
+                'confidence_tags': {},
+            },
+            'prediction_store_policy': prediction_store_policy,
+            'on_conflict': 'error',
+        },
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
+    # Arrange
+    data = {'test': 'data'}
+    path_flag = 'unknown'
+    confidence = 0.95
+    schema = 'test_schema'
+    table_name = 'test_table'
+    model = 'test_model'
+    last_timestamp = '2024-01-01'
+    model_name = 'test_model_name'
+    model_config = {'retention': '30'}
+    prediction_store_policy = 'erl:1'
+    # Act
+    result = await prediction_process.path_flag_handler(
+        data,
+        path_flag,
+        {
+            **metadata,
+            'schema': schema,
+            'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_config': model_config,
+            'opc_output_config': {'test': 'config'},
+            'pi_web_api_output_config': {
+                'endpoint': 'https://test-pi-server.com',
+                'prediction_tags': {},
+                'confidence_tags': {},
+            },
+            'prediction_store_policy': prediction_store_policy,
+        },
+        confidence,
+        last_timestamp,
+        '',
+    )
+
+    # Assert
+    assert result is False
+    workflow_mock.execute_activity_method.assert_not_called()
+    workflow_mock.execute_child_workflow.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process):
+    prediction_process.path_flag_handler = AsyncMock(return_value=False)
+    prediction_process.cleanup_prefixes = {'training_datasets/test'}
+
+    data_payload = MagicMock()
+    data_payload.cleanup_prefix.return_value = 'training_datasets/test'
+    data_payload.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    input_data = {
+        'metadata': metadata,
+        'data': data_payload,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'model_id': 1,
+        'input_filters': {'test': 'filter'},
+        'mlflow_transform_filters': {'test': 'filter'},
+        'mlflow_predict_filters': {'test': 'filter'},
+        'model_name': 'test_model_name',
+        'model_config': {'retention': '30'},
+        'path_priority': ['continue', 'repeat', 'stop'],
+        'opc_output_config': {'test': 'config'},
+        'pi_web_api_output_config': {'test': 'config'},
+        'prediction_store_policy': 'lts:1',
+    }
+
+    workflow_mock.execute_activity_method.side_effect = [
+        {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+        {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+        MagicMock(),
+    ]
+    workflow_mock.execute_local_activity_method.side_effect = [
+        ('continue', 0.95, 'ok'),
+        ('continue', 0.95, ''),
+        ('continue', 0.95, ''),
+        ('continue', 0.95, ''),
+    ]
+
+    await prediction_process.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_any_call(
+        Activities.cleanup_minio_objects_expired,
+        {**metadata, 'data': data_payload},
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )

tests/laborious/workflows/test_drift.py

@@ -0,0 +1,248 @@
+from unittest.mock import ANY, AsyncMock, call, patch
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.activities import Activities
+from laborious.workflows.drift import Drift
+
+
+@fixture
+def drift() -> Drift:
+    return Drift()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'drift',
+        'schedule_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+        'chunk_period': 'hour',
+    }
+
+    target_name = input_data['model_config']['target']
+
+    target_data = {'data': 'test_target_data'}
+    reference_data = {'data': 'test_reference_data'}
+    drift_data = {'drift': 'test_drift_data'}
+
+    workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
+    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Check start_local_activity_method calls
+    # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
+    expected_gathering_query = f"""
+            SELECT *
+            FROM "{input_data['schema']}"."{input_data['source_table_name']}"
+            WHERE
+                model_id = '{input_data['model_id']}' AND
+                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
+            ORDER BY timestamp ASC
+        """
+
+    workflow_mock.start_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                {
+                    **metadata,
+                    'query': expected_gathering_query,
+                    'datetime_columns': ['timestamp', 'created_at'],
+                    'orient': 'records',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.get_reference_data,
+                {
+                    **metadata,
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    # Assert - Check calculate_drift call
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_drift,
+        {
+            **metadata,
+            'target_data': target_data,
+            'reference_data': reference_data,
+            'model_name': input_data['model_name'],
+            'model_id': input_data['model_id'],
+            'target_name': target_name,
+            'drift_metrics': input_data['drift_metrics'],
+            'chunk_period': input_data['chunk_period'],
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    # Assert - Check export_data_to_postgres call
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.export_data_to_postgres,
+        {
+            **metadata,
+            'data': drift_data,
+            'schema': input_data['schema'],
+            'table_name': input_data['target_table_name'],
+            'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+    }
+
+    target_data = None
+    reference_data = {'data': 'test_reference_data'}
+
+    workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
+
+    workflow_mock.execute_activity_method = AsyncMock()
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Should not call calculate_drift or export
+    workflow_mock.execute_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+    }
+
+    target_name = input_data['model_config']['target']
+
+    target_data = {'data': 'test_target_data'}
+    reference_data = {'data': 'test_reference_data'}
+    drift_data = None
+
+    workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Should call calculate_drift but not export
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_drift,
+        {
+            **metadata,
+            'target_data': target_data,
+            'reference_data': reference_data,
+            'model_name': input_data['model_name'],
+            'model_id': input_data['model_id'],
+            'target_name': target_name,
+            'drift_metrics': input_data['drift_metrics'],
+            'chunk_period': input_data.get('chunk_period', 'min'),
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    workflow_mock.execute_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+        # chunk_period not provided, should default to 'min'
+    }
+
+    target_name = input_data['model_config']['target']
+
+    target_data = {'data': 'test_target_data'}
+    reference_data = {'data': 'test_reference_data'}
+    drift_data = {'drift': 'test_drift_data'}
+
+    workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
+    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Check calculate_drift call with default chunk_period
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_drift,
+        {
+            **metadata,
+            'target_data': target_data,
+            'reference_data': reference_data,
+            'model_name': input_data['model_name'],
+            'model_id': input_data['model_id'],
+            'target_name': target_name,
+            'drift_metrics': input_data['drift_metrics'],
+            'chunk_period': 'min',  # Default value
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )

tests/laborious/workflows/test_minimal_retrain.py

@@ -0,0 +1,327 @@
+from unittest.mock import ANY, AsyncMock, call, patch
+
+from pytest import fixture, mark
+
+from laborious.activities.activities import Activities
+from laborious.workflows.minimal_retrain import MinimalRetrain
+
+
+@fixture
+def minimal_retrain() -> MinimalRetrain:
+    return MinimalRetrain()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    storage_result = {
+        'last_timestamp': '2024-01-01 00:00:00+0000',
+        'status': {'success': True},
+        'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            storage_result,
+            {'success': True, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
+    )
+
+    await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_query_with_minio_offload,
+                {
+                    **metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.retrain_model,
+                {
+                    **metadata,
+                    'data': storage_result,
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.update_production_model,
+                {
+                    **metadata,
+                    'model_name': input_data['model_name'],
+                    'success': True,
+                    'experiment': 'test_experiment',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_retrain_report,
+                {
+                    **metadata,
+                    'experiment_response': {'success': True, 'experiment': 'test_experiment'},
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    'update_report': {
+                        'success': True,
+                        'version': 'test_version',
+                        'mlflow_run_id': 'test_mlflow_run_id',
+                        'mlflow_experiment_id': 'test_mlflow_experiment_id',
+                    },
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    storage_result = {
+        'last_timestamp': '2024-01-01 00:00:00+0000',
+        'status': {'success': True},
+        'data': {},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            storage_result,
+            {'success': True, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
+    )
+
+    from pytest import raises
+
+    with raises(ValueError, match='No data returned from query'):
+        await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.load_query_with_minio_offload,
+        {
+            **metadata,
+            'query': input_data['query'],
+            'datetime_columns': input_data.get('datetime_columns', []),
+            'model_name': input_data['model_name'],
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    workflow_mock.execute_local_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    storage_result = {
+        'last_timestamp': '2024-01-01 00:00:00+0000',
+        'status': {'success': True},
+        'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            storage_result,
+            {'success': False, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
+    )
+
+    await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_query_with_minio_offload,
+                {
+                    **metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.retrain_model,
+                {
+                    **metadata,
+                    'data': storage_result,
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_retrain_report,
+                {
+                    **metadata,
+                    'experiment_response': {'success': False, 'experiment': 'test_experiment'},
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    'update_report': {},
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/test_predictions_batch.py

@@ -0,0 +1,106 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+from pytest import fixture, mark
+
+from laborious.activities.activities import Activities
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+
+@fixture
+def predictions_batch() -> PredictionsBatch:
+    return PredictionsBatch()
+
+
+metadata = {
+    'model_id': 'test_model_id',
+    'model_name': 'test_model',
+    'workflow_name': 'predictions_batch',
+    'schedule_name': 'test_schedule',
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
+    activity_return = MagicMock()
+    workflow_mock.execute_activity_method.return_value = activity_return
+
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'query': 'SELECT * FROM test',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'opc_output_config': 'test_opc_output_config',
+        'pi_web_api_output_config': 'test_pi_web_api_output_config',
+        'datetime_columns': ['timestamp', 'created_at'],
+        'prediction_store_policy': 'erl:1',
+        'model_config': {'retention': '30'},
+    }
+
+    await predictions_batch.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_query_with_minio_offload,
+                {
+                    'metadata': metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    prediction_input = {
+        'metadata': {'metadata': metadata},
+        'data': activity_return,
+        'schema': input_data['schema'],
+        'table_name': input_data['table_name'],
+        'transform_table_name': input_data['transform_table_name'],
+        'model_id': input_data['model_id'],
+        'model_name': input_data['model_name'],
+        'input_filters': input_data.get(
+            'input_filters',
+            {
+                'EMPTY_DATA': {
+                    'POLICY': 'STOP',
+                    'CONFIG': {},
+                }
+            },
+        ),
+        'mlflow_transform_filters': input_data.get(
+            'mlflow_transform_filters',
+            {
+                'API_ERROR': {
+                    'POLICY': 'STOP',
+                    'CONFIG': {},
+                }
+            },
+        ),
+        'mlflow_predict_filters': input_data.get(
+            'mlflow_predict_filters',
+            {
+                'API_ERROR': {
+                    'POLICY': 'STOP',
+                    'CONFIG': {},
+                }
+            },
+        ),
+        'model_config': input_data.get('model_config', {}),
+        'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
+        'opc_output_config': input_data.get('opc_output_config', {}),
+        'on_conflict': input_data.get('on_conflict', 'error'),
+        'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
+        'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
+        'save_transform': input_data.get('save_transform', True),
+    }
+
+    workflow_mock.execute_child_workflow.assert_has_calls(
+        [call('subworkflow.prediction_process', prediction_input)]
+    )

tests/laborious/workflows/test_simple_metrics.py

@@ -0,0 +1,215 @@
+from unittest.mock import ANY, AsyncMock, call, patch
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.activities import Activities
+from laborious.workflows.simple_metrics import SimpleMetrics
+
+
+@fixture
+def simple_metrics() -> SimpleMetrics:
+    return SimpleMetrics()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'simple_metrics',
+        'schedule_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        'metrics': ['rmse', 'mse', 'mae', 'r2'],
+    }
+
+    target_data = {'data': 'test_target_data'}
+    simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
+
+    workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None])
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Check load_custom_query call
+    # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
+    expected_query = f"""
+            select p."timestamp", p.prediction, ld.value as "target"
+            from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
+            inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
+                on p."timestamp" = ld."timestamp"
+            where
+                p.model_id = '{input_data['model_id']}' and
+                p.prediction is not null and
+                ld.variable = '{input_data['model_config']['target']}' and
+                ld.value is not null and
+                p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
+            order by
+                p."timestamp" desc;
+        """
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                {
+                    **metadata,
+                    'query': expected_query,
+                    'datetime_columns': ['timestamp'],
+                    'orient': 'records',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': simple_metrics_data,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['target_table_name'],
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_simple_metrics,
+        {
+            **metadata,
+            'model_id': input_data['model_id'],
+            'target_data': target_data,
+            'metrics': input_data['metrics'],
+            'interval_minutes': input_data['interval_minutes'],
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        'metrics': ['rmse', 'mse'],
+    }
+
+    target_data = None
+
+    workflow_mock.execute_activity_method = AsyncMock(return_value=target_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Should not call calculate_simple_metrics or export
+    assert workflow_mock.execute_activity_method.call_count == 1
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        'metrics': ['rmse', 'mse'],
+    }
+
+    target_data = {'data': 'test_target_data'}
+    simple_metrics_data = None
+
+    workflow_mock.execute_activity_method = AsyncMock(return_value=target_data)
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Should call calculate_simple_metrics but not export
+    workflow_mock.execute_activity_method.assert_called_once()
+    workflow_mock.execute_local_activity_method.assert_called_once()
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        # metrics not provided, should default to ['rmse', 'mse', 'mae', 'r2']
+    }
+
+    target_data = {'data': 'test_target_data'}
+    simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
+
+    workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None])
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Check calculate_simple_metrics call with default metrics
+    workflow_mock.execute_activity_method.assert_any_call(
+        Activities.load_custom_query,
+        ANY,
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_simple_metrics,
+        {
+            **metadata,
+            'model_id': input_data['model_id'],
+            'target_data': target_data,
+            'metrics': ['rmse', 'mse', 'mae', 'r2'],  # Default value
+            'interval_minutes': input_data['interval_minutes'],
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )