Code import - branch release/SIENTIAPDE-1646

tests/laborious/activities/test_activities.py

@@ -0,0 +1,299 @@
+from unittest.mock import ANY, MagicMock, patch
+
+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_repository = MagicMock()
+    plugin_store = MagicMock()
+
+    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,
+        plugin_store=plugin_store,
+        minio_config=minio_config,
+        opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
+        logger=logger,
+        notification_handler=notification_handler,
+        mlflow_repository=mlflow_repository,
+    )
+
+    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_repository=mlflow_repository,
+        plugin_store=plugin_store,
+        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'],
+    )
+
+
+@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')
+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 = MagicMock()
+    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_repository = MagicMock()
+    plugin_store = MagicMock()
+
+    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,
+        plugin_store=plugin_store,
+        minio_config=minio_config,
+        opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
+        logger=logger,
+        notification_handler=notification_handler,
+        mlflow_repository=mlflow_repository,
+    )
+
+    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()
+
+
+@patch('laborious.activities.activities.SientiaMLflowRepository')
+@patch('laborious.activities.activities.build_mlflow_config')
+@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___builds_mlflow_repository_when_not_provided(
+    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,
+    mock_build_mlflow_config,
+    mock_mlflow_repository_cls,
+):
+    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,
+    }
+    opc_config = {'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, 'group_id': 'test'}
+    pi_web_api_config = {'base_url': 'https://pi', 'auth_type': 'bearer', 'auth_token': 'token'}
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    plugin_store = MagicMock()
+    mock_build_mlflow_config.return_value = {
+        'url': 'http://mlflow:80',
+        'username': 'u',
+        'password': 'p',
+    }
+
+    Activities(
+        postgres_config=postgres_config,
+        plugin_store=plugin_store,
+        minio_config=minio_config,
+        opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
+        logger=logger,
+        notification_handler=notification_handler,
+    )
+
+    mock_build_mlflow_config.assert_called_once()
+    mock_mlflow_repository_cls.assert_called_once_with(
+        host='http://mlflow:80',
+        username='u',
+        password='p',
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )

tests/laborious/activities/test_api.py

@@ -0,0 +1,503 @@
+from unittest.mock import ANY, MagicMock, call, patch
+
+from pytest import fixture
+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=MagicMock(),
+    )
+    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=MagicMock(),
+    )
+    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=MagicMock(),
+    )
+
+    assert api.pi_web_api_client is not None
+
+
+@fixture
+@patch('laborious.activities.api.PIWebAPIClient')
+def api(mock_pi_web_api_client):
+    mock_client = MagicMock()
+    mock_client.write_value = MagicMock()
+    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=MagicMock(),
+    )
+    api_instance.send_notification = MagicMock()
+    api_instance.info = MagicMock()
+    api_instance.error = MagicMock()
+    api_instance.emit_metric_sync = MagicMock()
+    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
+
+
+@patch('laborious.activities.api.DataFrame')
+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 = 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'],
+    }
+
+
+@patch('laborious.activities.api.DataFrame')
+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 = api.write_pi_web_api_data(base_input_data)
+
+    api.send_notification.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
+
+
+@patch('laborious.activities.api.DataFrame')
+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 = api.write_pi_web_api_data(base_input_data)
+
+    api.send_notification.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
+
+
+@patch('laborious.activities.api.DataFrame')
+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 = 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'],
+    }
+
+
+@patch('laborious.activities.api.DataFrame')
+def test_write_pi_web_api_data_updates_confidence_and_comments(
+    mock_dataframe, api, base_input_data
+):
+    mock_dataframe.return_value = _create_mock_dataframe()
+    api.pi_web_api_client.write_value.side_effect = [
+        [{'WebId': 'web_id_1', 'Errors': []}],
+        [{'WebId': 'web_id_2', 'Errors': []}],
+    ]
+    with patch.object(
+        api,
+        'process_pi_web_api_response',
+        new=MagicMock(side_effect=[(0.33, 'PI warning'), (0, '')]),
+    ) as process_mock:
+        result = api.write_pi_web_api_data(base_input_data)
+
+    assert process_mock.call_count == 2
+    assert result is not None
+
+
+@patch('laborious.activities.api.SientiaMonitoring.shutdown')
+def test_close(mock_shutdown, api):
+    api.close()
+
+    api.pi_web_api_client.close.assert_called_once()
+    mock_shutdown.assert_called_once_with(api)
+
+
+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 = 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_sync.call_count == 2
+    # Verify that emit_metric_sync was called with correct tags structure
+    call_args_list = api.emit_metric_sync.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']
+
+
+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 = 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_sync.call_count == 2
+
+
+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 = 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.assert_called_once()
+    call_args = api.send_notification.call_args
+    assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
+    assert call_args.kwargs['level'] == NotificationLevel.ERROR
+
+
+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 = 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'])
+
+
+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 = 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,856 @@
+from datetime import datetime
+from unittest.mock import ANY, MagicMock, patch
+
+import numpy as np
+import pandas as pd
+from pytest import fixture, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.constants import 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.MinioRepository')
+def test___init__(mock_minio_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = MagicMock()
+    mlflow_repo = MagicMock()
+    plugin_store = MagicMock()
+
+    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_repository=mlflow_repo,
+        plugin_store=plugin_store,
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    assert mlflow.mlflow_repository is mlflow_repo
+    assert mlflow.plugin_store is plugin_store
+
+    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.MinioRepository')
+def mlflow(mock_minio_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = MagicMock()
+    mlflow_repo = MagicMock()
+    plugin_store = MagicMock()
+
+    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_repository=mlflow_repo,
+        plugin_store=plugin_store,
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    mlflow.minio_repository = MagicMock()
+
+    mlflow.send_notification = MagicMock()
+    mlflow.emit_metric = MagicMock()
+    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',
+    },
+}
+
+
+def test_detect_and_parse_datetime_index_empty(mlflow):
+    df = pd.DataFrame()
+    out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
+    assert out.empty
+
+
+def test_detect_and_parse_datetime_index_mixed_types_error(mlflow):
+    idx = pd.Index([pd.Timestamp('2020-01-01', tz='UTC'), 'x'])
+    df = pd.DataFrame({'a': [1, 2]}, index=idx)
+    with raises(ValueError):
+        mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
+
+
+def test_detect_and_parse_datetime_index_invalid_string_error(mlflow):
+    idx = pd.Index(['bad-format'])
+    df = pd.DataFrame({'a': [1]}, index=idx)
+    with raises(ValueError):
+        mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
+
+
+def test_detect_and_parse_datetime_index_unsupported_type_error(mlflow):
+    idx = pd.Index([pd.Period('2020-01', freq='M')])
+    df = pd.DataFrame({'a': [1]}, index=idx)
+    with raises(ValueError):
+        mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
+
+
+def test_detect_and_parse_datetime_index_datetime_success(mlflow):
+    idx = pd.Index([datetime(2020, 1, 1, 0, 0, 0)])
+    df = pd.DataFrame({'a': [1]}, index=idx)
+    out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
+    assert out.index[0].endswith('+0000')
+
+
+def test_detect_and_parse_datetime_index_timestamp_with_tz_success(mlflow):
+    idx = pd.DatetimeIndex([pd.Timestamp('2020-01-01 00:00:00', tz='UTC')])
+    df = pd.DataFrame({'a': [1]}, index=idx)
+    out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
+    assert out.index[0].endswith('+0000')
+
+
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=MagicMock,
+)
+def test_request_transform_success(mock_from_dataframe, mlflow):
+    ts = pd.Timestamp('2020-01-01', tz='UTC')
+    raw = pd.DataFrame(
+        {
+            'variable': ['v1', 'v1'],
+            'timestamp': [ts, ts],
+            'value': [1.0, 2.0],
+            'created_at': [ts, ts],
+        }
+    )
+    pivoted = raw.sort_values('created_at', ascending=False).drop_duplicates(
+        subset=['variable', 'timestamp'], keep='first'
+    )
+    pivoted = pivoted.pivot(index='timestamp', columns='variable', values='value')
+    pivoted = pivoted.fillna(np.nan)
+    pivoted.columns.name = None
+    pivoted.index.name = None
+    pivoted['timestamp'] = pivoted.index
+
+    out_idx = pd.Index([ts.strftime(DATETIME_FORMAT_WITH_TZ)], name=None)
+    out_df = pd.DataFrame({'v1': [1.0]}, index=out_idx)
+    wrapper = MagicMock()
+    wrapper.transform.return_value = (out_df, {'meta': True})
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    response_data = mlflow.request_transform(input_data)
+
+    mlflow.mlflow_repository.get_cached_model.assert_called_once_with(
+        model_name='test_model',
+        alias='production',
+        retention_minutes=0,
+        metadata=metadata['metadata'],
+    )
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
+
+
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=MagicMock,
+)
+def test_request_transform_success_without_transform_meta(mock_from_dataframe, mlflow):
+    ts = pd.Timestamp('2020-01-01', tz='UTC')
+    raw = pd.DataFrame(
+        {
+            'variable': ['v1'],
+            'timestamp': [ts],
+            'value': [1.0],
+            'created_at': [ts],
+        }
+    )
+    out_idx = pd.Index([ts.strftime(DATETIME_FORMAT_WITH_TZ)], name=None)
+    out_df = pd.DataFrame({'v1': [1.0]}, index=out_idx)
+    wrapper = MagicMock()
+    wrapper.transform.return_value = (out_df, {})
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    mlflow.request_transform(input_data)
+    mock_from_dataframe.assert_called_once()
+
+
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=MagicMock,
+)
+def test_request_transform_failure(mock_from_dataframe, mlflow):
+    mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('boom')
+
+    data_mock = MagicMock()
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    data_mock.sort_values.return_value = data_mock
+    data_mock.drop_duplicates.return_value = data_mock
+    data_mock.pivot.return_value = data_mock
+
+    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={'success': False, 'content': ANY},
+        workflow_metadata=metadata['metadata'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+
+
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=MagicMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
+    wrapper = MagicMock()
+    pred_df = MagicMock()
+    wrapper.predict.return_value = (pred_df, {})
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+
+    data_mock = MagicMock()
+    data_mock.index = pd.DatetimeIndex([pd.Timestamp('2020-01-01', tz='UTC')])
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    pred_df.columns = MagicMock()
+    pred_df.__setitem__ = MagicMock()
+
+    response_data = mlflow.request_predict(input_data)
+
+    data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
+    mock_to_datetime.assert_called()
+    mlflow.mlflow_repository.get_cached_model.assert_called_once()
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
+
+
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=MagicMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+def test_request_predict_success_dataframe_and_meta(mock_to_datetime, mock_from_dataframe, mlflow):
+    mock_to_datetime.side_effect = lambda x, **kwargs: x
+    wrapper = MagicMock()
+    pred_df = pd.DataFrame({'raw': [0.3]})
+    wrapper.predict.return_value = (pred_df, {'m': 1})
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+
+    data_mock = MagicMock()
+    data_mock.index = pd.DatetimeIndex([pd.Timestamp('2020-01-01', tz='UTC')])
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    mlflow.request_predict(input_data)
+
+    assert list(pred_df.columns) == ['prediction', 'response_time']
+    mlflow.info.assert_any_call("Wrapper predict metadata: {'m': 1}", metadata['metadata'])
+    mock_from_dataframe.assert_called_once()
+
+
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=MagicMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
+    mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('predict boom')
+
+    data_mock = MagicMock()
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    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={'success': False, 'content': ANY},
+        workflow_metadata=metadata['metadata'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+
+
+@patch('laborious.activities.mlflow.mlflow.log_artifact')
+@patch('laborious.activities.mlflow.tempfile.mkdtemp')
+@patch('laborious.activities.mlflow.rmtree')
+@patch('laborious.activities.mlflow.to_datetime')
+def test_retrain_model_success_data_success_retrain(
+    mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
+):
+    mock_mkdtemp.return_value = 'tmp'
+    mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
+
+    mv_alias = MagicMock()
+    mv_alias.run_id = 'source-run'
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
+
+    wrapper = MagicMock()
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+
+    mock_cm = MagicMock()
+    mock_cm.__enter__.return_value = MagicMock(run_id='new-run', experiment_id='exp-1')
+    mock_cm.__exit__.return_value = False
+    mlflow.mlflow_repository.start_run.return_value = mock_cm
+
+    ts = pd.Timestamp('2020-01-01', tz='UTC')
+    raw_data = pd.DataFrame(
+        {
+            'variable': ['target', 'f1'],
+            'timestamp': [ts, ts],
+            'value': [1.0, 2.0],
+        }
+    )
+    pivoted_index = pd.DatetimeIndex([ts])
+    wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0])
+
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw_data)
+
+    response = mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+            },
+        }
+    )
+
+    wrapper.retrain.assert_called_once()
+    wrapper.store_model.assert_called_once_with(name='test_model')
+    assert response['success'] is True
+    assert response['experiment']['run_id'] == 'new-run'
+
+
+@patch('laborious.activities.mlflow.mlflow.log_artifact')
+@patch('laborious.activities.mlflow.tempfile.mkdtemp')
+@patch('laborious.activities.mlflow.rmtree')
+@patch('laborious.activities.mlflow.to_datetime')
+def test_retrain_model_success_with_payload_data(
+    mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
+):
+    mock_mkdtemp.return_value = 'tmp'
+    mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
+    mv_alias = MagicMock(run_id='src')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
+    wrapper = MagicMock()
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+    mock_cm = MagicMock()
+    mock_cm.__enter__.return_value = MagicMock(run_id='r', experiment_id='e')
+    mock_cm.__exit__.return_value = False
+    mlflow.mlflow_repository.start_run.return_value = mock_cm
+
+    ts = pd.Timestamp('2020-01-01', tz='UTC')
+    raw_data = pd.DataFrame(
+        {
+            'variable': ['target', 'f1', 'target', 'f1'],
+            'timestamp': [ts, ts, ts + pd.Timedelta(hours=1), ts + pd.Timedelta(hours=1)],
+            'value': [1.0, 2.0, 3.0, 4.0],
+            'created_at': [ts, ts, ts + pd.Timedelta(hours=1), ts + pd.Timedelta(hours=1)],
+        }
+    )
+    pivoted_index = pd.DatetimeIndex([ts, ts + pd.Timedelta(hours=1)])
+    wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0, 3.0])
+
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw_data)
+
+    response = mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {'target': 'target'},
+        }
+    )
+
+    assert response['success'] is True
+
+
+@patch('laborious.activities.mlflow.mlflow.log_artifact')
+@patch('laborious.activities.mlflow.tempfile.mkdtemp')
+@patch('laborious.activities.mlflow.rmtree')
+@patch('laborious.activities.mlflow.to_datetime')
+def test_retrain_model_always_uses_retrain_even_with_full_retrain_flag(
+    mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
+):
+    mock_mkdtemp.return_value = 'tmp'
+    mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
+    mv_alias = MagicMock(run_id='src')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
+    wrapper = MagicMock()
+    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
+    mock_cm = MagicMock()
+    mock_cm.__enter__.return_value = MagicMock(run_id='r', experiment_id='e')
+    mock_cm.__exit__.return_value = False
+    mlflow.mlflow_repository.start_run.return_value = mock_cm
+
+    ts = pd.Timestamp('2020-01-01', tz='UTC')
+    ts_next = ts + pd.Timedelta(days=1)
+    raw_data = pd.DataFrame(
+        {
+            'variable': ['target', 'f1', 'target', 'f1'],
+            'timestamp': [ts, ts, ts_next, ts_next],
+            'value': [1.0, 2.0, 3.0, 4.0],
+        }
+    )
+    pivoted_index = pd.DatetimeIndex([ts, ts_next])
+    wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0, 3.0])
+
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw_data)
+
+    response = mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {'target': 'target', 'full_retrain': True, 'validation_fraction': 0.5},
+        }
+    )
+
+    wrapper.retrain.assert_called_once()
+    wrapper.train.assert_not_called()
+    assert response['success'] is True
+
+
+@patch('laborious.activities.mlflow.to_datetime')
+def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
+    mv_alias = MagicMock(run_id='src')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
+    mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('retrain failed')
+
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
+    raw_data.__getitem__.return_value.max.return_value = 'tsmax'
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw_data)
+
+    pivoted = MagicMock()
+    raw_data.sort_values.return_value = raw_data
+    raw_data.drop_duplicates.return_value = raw_data
+    raw_data.pivot.return_value = pivoted
+    pivoted.fillna = MagicMock()
+    pivoted.columns.name = None
+    pivoted.index = MagicMock()
+    pivoted.__setitem__ = MagicMock()
+
+    response = mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {'target': 'target'},
+        }
+    )
+
+    assert response['success'] is False
+    assert 'retrain failed' in response['message']
+
+
+def test_retrain_model_data_error(mlflow):
+    response = mlflow.retrain_model(
+        {
+            **metadata,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+            },
+        }
+    )
+
+    assert response['success'] is False
+    assert 'data' in response['message'].lower() or 'loading' in response['message'].lower()
+
+
+def test_retrain_model_missing_target(mlflow):
+    ts = pd.Timestamp('2020-01-01', tz='UTC')
+    raw_data = pd.DataFrame(
+        {
+            'variable': ['f1', 'f2'],
+            'timestamp': [ts, ts],
+            'value': [1.0, 2.0],
+        }
+    )
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=raw_data)
+
+    response = mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {},
+        }
+    )
+
+    assert response['success'] is False
+    assert 'target' in response['message']
+
+
+def test_retrain_model_data_error_no_minio_repository(mlflow):
+    mlflow.minio_repository = None
+
+    with raises(ValueError) as e:
+        mlflow.retrain_model(
+            {
+                **metadata,
+                'object_key': 'test_object_key',
+                'model_name': 'test_model',
+                'model_config': {
+                    'target': 'target',
+                },
+            }
+        )
+
+    assert str(e.value) == 'Minio repository not initialized'
+
+
+def test_update_production_model(mlflow):
+    mlflow.mlflow_repository._client.search_model_versions.return_value = [
+        MagicMock(version='3', run_id='run-x'),
+        MagicMock(version='2', run_id='run-x'),
+    ]
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 1,
+        'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
+        'timestamp': 2,
+        'status': 'success',
+    }
+
+    response = mlflow.update_production_model(input_data)
+
+    mlflow.mlflow_repository.promote_to_alias.assert_called_once_with(
+        model_name='test_model',
+        version='3',
+        alias='production',
+        metadata=metadata['metadata'],
+    )
+
+    assert response['model_name'] == 'test_model'
+    assert response['version'] == '3'
+
+
+def test_update_production_model_error(mlflow):
+    mlflow.mlflow_repository._client.search_model_versions.return_value = []
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 1,
+        'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
+        'timestamp': 2,
+        'status': 'success',
+    }
+
+    try:
+        mlflow.update_production_model(input_data)
+    except Exception:
+        mlflow.send_notification.assert_called_once_with(
+            metadata=metadata['metadata'],
+            notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
+            message=ANY,
+            block='update_production_model',
+            level=NotificationLevel.ERROR,
+            attachment_content=ANY,
+        )
+    else:
+        raise AssertionError('Expected exception')
+
+
+def _artifact_file_info(path: str) -> MagicMock:
+    file_info = MagicMock()
+    file_info.path = path
+    return file_info
+
+
+def _retrain_prediction_frame(
+    index: pd.DatetimeIndex, target: str, values: list[float]
+) -> pd.DataFrame:
+    return pd.DataFrame({target: values}, index=index)
+
+
+@patch('laborious.activities.mlflow.to_datetime')
+def test_get_reference_data_success(mock_to_datetime, mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('evaluation_data.csv'),
+    ]
+
+    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},
+    ]
+
+    with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+        with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+            with patch('laborious.activities.mlflow.rmtree'):
+                with patch.object(
+                    mlflow,
+                    '_find_downloaded_csv',
+                    return_value=MagicMock(),
+                ):
+                    result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='evaluation_data.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    mock_reference_data.to_dict.assert_called_once_with(orient='records')
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_not_found(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.mlflow_repository._client.get_model_version_by_alias.side_effect = Exception('missing')
+
+    result = mlflow.get_reference_data(input_data)
+
+    mlflow.warning.assert_called()
+    assert result is None
+
+
+def test_get_reference_data_only_test_data_csv(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('test_data.csv'),
+    ]
+
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
+
+    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
+        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+                with patch('laborious.activities.mlflow.rmtree'):
+                    with patch.object(
+                        mlflow,
+                        '_find_downloaded_csv',
+                        return_value=MagicMock(),
+                    ):
+                        result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='test_data.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_prefers_evaluation_data_when_both_listed(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('test_data.csv'),
+        _artifact_file_info('evaluation_data.csv'),
+    ]
+
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
+
+    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
+        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
+            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
+                with patch('laborious.activities.mlflow.rmtree'):
+                    with patch.object(
+                        mlflow,
+                        '_find_downloaded_csv',
+                        return_value=MagicMock(),
+                    ):
+                        result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
+        run_id='run1',
+        artifact_path='evaluation_data.csv',
+        dst_path='/t',
+        metadata=metadata['metadata'],
+    )
+    assert result == mock_reference_data.to_dict.return_value
+
+
+def test_get_reference_data_no_candidate_artifacts_returns_none(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('other_artifact.csv'),
+    ]
+
+    result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_not_called()
+    mlflow.warning.assert_called()
+    assert result is None
+
+
+def test_get_reference_data_missing_csv_file_returns_none(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('evaluation_data.csv'),
+    ]
+
+    with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='tmp'):
+        with patch('laborious.activities.mlflow.rmtree'):
+            with patch.object(mlflow, '_find_downloaded_csv', return_value=None):
+                result = mlflow.get_reference_data(input_data)
+
+    mlflow.mlflow_repository.download_artifacts.assert_called_once()
+    mlflow.warning.assert_called()
+    assert result is None
+
+
+def test_get_reference_data_exception(mlflow):
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mv = MagicMock(run_id='run1')
+    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
+    mlflow.mlflow_repository._client.list_artifacts.return_value = [
+        _artifact_file_info('evaluation_data.csv'),
+    ]
+    mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')
+
+    result = mlflow.get_reference_data(input_data)
+
+    assert result is None

tests/laborious/activities/test_model_metrics.py

@@ -0,0 +1,954 @@
+from unittest.mock import ANY, MagicMock, patch
+
+from pandas import DataFrame, Timestamp
+from pytest import fixture, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+from sientia_model.analytics.drift_analysis import DriftInsufficientDataError
+
+from laborious.activities.model_metrics import ModelMetrics
+
+
+@fixture
+def model_metrics_activity():
+    model_metrics = ModelMetrics(
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=MagicMock(),
+    )
+    model_metrics.error = MagicMock()
+    model_metrics.debug = MagicMock()
+    model_metrics.info = MagicMock()
+    model_metrics.warning = MagicMock()
+    model_metrics.critical = MagicMock()
+    model_metrics.send_notification = MagicMock()
+    model_metrics.emit_metric_sync = MagicMock()
+    model_metrics.get_core_labels = MagicMock(
+        return_value={
+            'pod_id': 'test_pod',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+        }
+    )
+    model_metrics.observe_lag_sync = MagicMock()
+    model_metrics.pod_id = 'test_pod'
+    return model_metrics
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': None,
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'invalid',
+    }
+
+    # Act & Assert
+    try:
+        model_metrics_activity.calculate_drift(input_data)
+    except ValueError as e:
+        assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
+        model_metrics_activity.error.assert_called_once_with(
+            'Invalid chunk period: invalid', metadata['metadata']
+        )
+    else:
+        raise AssertionError('Expected ValueError')
+
+
+def _sample_drift_metrics_df(ts: Timestamp) -> DataFrame:
+    """Minimal analyzer-shaped dataframe (univariate row + columns the activity expects)."""
+    return DataFrame(
+        {
+            'timestamp': [ts],
+            'feature': ['feature1'],
+            'method': ['ks_test'],
+            'value': [0.5],
+            'alert': [False],
+            'chunk_index': [0],
+            'chunk_start_date': [ts],
+            'chunk_end_date': [ts],
+            'threshold': [0.1],
+            'drift_type': ['univariate'],
+        }
+    )
+
+
+def test_calculate_drift_with_reference_data(model_metrics_activity):
+    ts = Timestamp('2023-05-26 11:12:27')
+    drift_df = _sample_drift_metrics_df(ts)
+    model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict('list'),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    result = model_metrics_activity.calculate_drift(input_data)
+
+    expected_timestamp = ts.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
+    assert result == [
+        {
+            'timestamp': expected_timestamp,
+            'feature': 'feature1',
+            'method': 'ks_test',
+            'value': 0.5,
+            'alert': False,
+            'chunk_index': 0,
+            'chunk_start_date': ts.isoformat(),
+            'chunk_end_date': ts.isoformat(),
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+    model_metrics_activity.info.assert_called()
+    model_metrics_activity.get_drift_metrics.assert_called_once()
+
+
+def test_calculate_drift_without_reference_data(model_metrics_activity):
+    # Ten rows so int(len * 0.3) >= 1 for the built-in reference slice.
+    ts_last = Timestamp('2023-05-26 11:12:36')
+    drift_df = _sample_drift_metrics_df(ts_last)
+    model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
+
+    timestamps = [f'2023-05-26 11:12:{27 + i:02d}' for i in range(10)]
+    target_data_dict = {
+        'timestamp': timestamps,
+        'variable': ['feature1'] * 10,
+        'value': [float(i) for i in range(10)],
+    }
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': None,
+        'target_data': target_data_dict,
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 's',
+    }
+
+    result = model_metrics_activity.calculate_drift(input_data)
+
+    expected_timestamp = ts_last.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
+    assert result == [
+        {
+            'timestamp': expected_timestamp,
+            'feature': 'feature1',
+            'method': 'ks_test',
+            'value': 0.5,
+            'alert': False,
+            'chunk_index': 0,
+            'chunk_start_date': ts_last.isoformat(),
+            'chunk_end_date': ts_last.isoformat(),
+            'model_id': 'test_model_id',
+            'accurate': False,
+        }
+    ]
+    model_metrics_activity.warning.assert_called()
+    model_metrics_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
+        message='Using 30% first rows of target data as reference data',
+        block='model_metrics',
+        level=NotificationLevel.WARNING,
+        attachment_content=ANY,
+    )
+
+
+def test_calculate_drift_empty_drift_df(model_metrics_activity):
+    """Empty analyzer merge yields no rows and no insufficient-data alert (lib owns that failure mode)."""
+    ts = Timestamp('2023-05-26 11:12:27')
+    empty_df = _sample_drift_metrics_df(ts).iloc[0:0]
+    model_metrics_activity.get_drift_metrics = MagicMock(return_value=empty_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    result = model_metrics_activity.calculate_drift(input_data)
+    assert result == []
+    model_metrics_activity.send_notification.assert_not_called()
+
+
+def test_calculate_drift_empty_after_timestamp_filter(model_metrics_activity):
+    """Rows dropped by target-window alignment yield an empty export list, not an insufficient-data error."""
+    drift_df = _sample_drift_metrics_df(Timestamp('2020-01-01'))
+    model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    result = model_metrics_activity.calculate_drift(input_data)
+    assert result == []
+    model_metrics_activity.send_notification.assert_not_called()
+
+
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+def test_calculate_drift_drift_insufficient_data_error_from_lib(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    """``DriftInsufficientDataError`` maps to MODEL_METRICS_DRIFT_INSUFFICIENT_DATA, not GET error."""
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+
+    lib_msg = (
+        '[MODEL_METRICS_DRIFT_INSUFFICIENT_DATA] Drift analysis produced no time chunks '
+        "(chunk_period='min', analysis_rows=1)."
+    )
+    model_metrics_activity.get_drift_metrics = MagicMock(
+        side_effect=DriftInsufficientDataError(lib_msg, analysis_rows=1)
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    with raises(DriftInsufficientDataError, match='MODEL_METRICS_DRIFT_INSUFFICIENT_DATA'):
+        model_metrics_activity.calculate_drift(input_data)
+
+    model_metrics_activity.error.assert_called_once_with(lib_msg, metadata['metadata'])
+    model_metrics_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='MODEL_METRICS_DRIFT_INSUFFICIENT_DATA',
+        message=lib_msg,
+        block='model_metrics',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+def test_calculate_drift_success_min(model_metrics_activity):
+    ts = Timestamp('2023-05-26 11:12:27')
+    drift_df = _sample_drift_metrics_df(ts)
+    model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict('list'),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    result = model_metrics_activity.calculate_drift(input_data)
+
+    expected_timestamp = ts.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
+    assert result == [
+        {
+            'timestamp': expected_timestamp,
+            'feature': 'feature1',
+            'method': 'ks_test',
+            'value': 0.5,
+            'alert': False,
+            'chunk_index': 0,
+            'chunk_start_date': ts.isoformat(),
+            'chunk_end_date': ts.isoformat(),
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+    model_metrics_activity.info.assert_called()
+    model_metrics_activity.get_drift_metrics.assert_called_once()
+
+
+def test_calculate_drift_success_s(model_metrics_activity):
+    ts = Timestamp('2023-05-26 11:12:27')
+    drift_df = _sample_drift_metrics_df(ts)
+    model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict('list'),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 's',
+    }
+
+    result = model_metrics_activity.calculate_drift(input_data)
+
+    expected_timestamp = ts.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
+    assert result == [
+        {
+            'timestamp': expected_timestamp,
+            'feature': 'feature1',
+            'method': 'ks_test',
+            'value': 0.5,
+            'alert': False,
+            'chunk_index': 0,
+            'chunk_start_date': ts.isoformat(),
+            'chunk_end_date': ts.isoformat(),
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+    model_metrics_activity.info.assert_called()
+    model_metrics_activity.get_drift_metrics.assert_called_once()
+
+
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+def test_calculate_drift_get_drift_metrics_error(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+
+    model_metrics_activity.get_drift_metrics = MagicMock(
+        side_effect=Exception('Get drift metrics error')
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    # Act / Assert
+    with raises(Exception, match='Get drift metrics error'):
+        model_metrics_activity.calculate_drift(input_data)
+
+    model_metrics_activity.error.assert_called_once_with(
+        'Error getting drift metrics: Get drift metrics error', metadata['metadata']
+    )
+    model_metrics_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
+        message='Error getting drift metrics: Get drift metrics error',
+        block='model_metrics',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+@patch('laborious.activities.model_metrics.to_datetime')
+@patch('laborious.activities.model_metrics.time.time')
+@patch('laborious.activities.model_metrics.DriftAnalysis')
+@patch('laborious.activities.model_metrics.metrics')
+def test_get_drift_metrics_success(
+    mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
+):
+    # Arrange
+    mock_time.return_value = 1000.0
+
+    mock_drift_df = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'method': ['ks_test'],
+            'value': [0.5],
+            'feature': ['feature1'],
+        }
+    )
+
+    mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    reference_columns = reference_data.drop(
+        columns=['target', 'timestamp'], errors='ignore'
+    ).columns
+
+    # Act
+    result = model_metrics_activity.get_drift_metrics(
+        reference_data=reference_data,
+        target_data=target_data,
+        target_name='target',
+        reference_columns=reference_columns,
+        drift_metrics=['ks_test'],
+        chunk_period='min',
+        metadata=metadata['metadata'],
+    )
+
+    # Assert
+    assert isinstance(result, DataFrame)
+    model_metrics_activity.debug.assert_called()
+    model_metrics_activity.observe_lag_sync.assert_called()
+    model_metrics_activity.emit_metric_sync.assert_called()
+
+
+@patch('laborious.activities.model_metrics.to_datetime')
+@patch('laborious.activities.model_metrics.time.time')
+@patch('laborious.activities.model_metrics.DriftAnalysis')
+@patch('laborious.activities.model_metrics.metrics')
+def test_get_drift_metrics_univariate_error(
+    mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
+):
+    # Arrange
+    mock_time.return_value = 1000.0
+
+    mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
+        'Univariate drift error'
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    reference_columns = reference_data.drop(
+        columns=['target', 'timestamp'], errors='ignore'
+    ).columns
+
+    # Act & Assert
+    try:
+        model_metrics_activity.get_drift_metrics(
+            reference_data=reference_data,
+            target_data=target_data,
+            target_name='target',
+            reference_columns=reference_columns,
+            drift_metrics=['ks_test'],
+            chunk_period='min',
+            metadata=metadata['metadata'],
+        )
+    except Exception as e:
+        assert str(e) == 'Univariate drift error'
+        model_metrics_activity.error.assert_called_once_with(
+            'Error detecting univariate drift: Univariate drift error', metadata['metadata']
+        )
+        model_metrics_activity.emit_metric_sync.assert_called_with(
+            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
+        )
+    else:
+        raise AssertionError('Expected Exception')
+
+
+@patch('laborious.activities.model_metrics.to_datetime')
+@patch('laborious.activities.model_metrics.time.time')
+@patch('laborious.activities.model_metrics.DriftAnalysis')
+@patch('laborious.activities.model_metrics.metrics')
+def test_get_drift_metrics_multivariate_error(
+    mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
+):
+    mock_time.return_value = 1000.0
+
+    mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.detect_multivariate_drift.side_effect = Exception(
+        'Multivariate drift error'
+    )
+
+    reference_data = DataFrame(
+        {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
+    )
+    target_data = DataFrame(
+        {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
+    )
+    reference_columns = reference_data.drop(
+        columns=['target', 'timestamp'], errors='ignore'
+    ).columns
+
+    try:
+        model_metrics_activity.get_drift_metrics(
+            reference_data=reference_data,
+            target_data=target_data,
+            target_name='target',
+            reference_columns=reference_columns,
+            drift_metrics=['ks_test'],
+            chunk_period='min',
+            metadata=metadata['metadata'],
+        )
+    except Exception as e:
+        assert str(e) == 'Multivariate drift error'
+        model_metrics_activity.error.assert_called_once_with(
+            'Error detecting multivariate drift: Multivariate drift error', metadata['metadata']
+        )
+        model_metrics_activity.emit_metric_sync.assert_called_with(
+            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
+        )
+    else:
+        raise AssertionError('Expected Exception')
+
+
+@patch('laborious.activities.model_metrics.to_datetime')
+@patch('laborious.activities.model_metrics.time.time')
+@patch('laborious.activities.model_metrics.DriftAnalysis')
+@patch('laborious.activities.model_metrics.metrics')
+def test_get_drift_metrics_dataframe_error(
+    mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
+):
+    mock_time.return_value = 1000.0
+
+    mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.get_drift_metrics_dataframe.side_effect = Exception(
+        'Dataframe error'
+    )
+
+    reference_data = DataFrame(
+        {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
+    )
+    target_data = DataFrame(
+        {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
+    )
+    reference_columns = reference_data.drop(
+        columns=['target', 'timestamp'], errors='ignore'
+    ).columns
+
+    try:
+        model_metrics_activity.get_drift_metrics(
+            reference_data=reference_data,
+            target_data=target_data,
+            target_name='target',
+            reference_columns=reference_columns,
+            drift_metrics=['ks_test'],
+            chunk_period='min',
+            metadata=metadata['metadata'],
+        )
+    except Exception as e:
+        assert str(e) == 'Dataframe error'
+        model_metrics_activity.error.assert_called_once_with(
+            'Error building drift metrics dataframe: Dataframe error', metadata['metadata']
+        )
+        model_metrics_activity.emit_metric_sync.assert_called_with(
+            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
+        )
+    else:
+        raise AssertionError('Expected Exception')
+
+
+def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+            'target': [1.0, 2.0, 3.0],
+            'prediction': [1.1, 2.1, 2.9],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['rmse', 'mse', 'mae', 'r2'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 4
+    assert 'rmse' in result['metric'].values
+    assert 'mse' in result['metric'].values
+    assert 'mae' in result['metric'].values
+    assert 'r2' in result['metric'].values
+    assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
+    assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
+    assert all(data_size == 3 for data_size in result['data_size'].values)
+    assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
+        metadata['metadata'],
+    )
+    model_metrics_activity.debug.assert_called_once()
+
+
+def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['rmse'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'rmse'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
+    )
+
+
+def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['mse'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'mse'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
+    )
+
+
+def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['mae'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'mae'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
+    )
+
+
+def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['r2'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'r2'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
+    )
+
+
+def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
+    # Arrange
+    # All target values are the same, so ss_tot will be 0
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 1.0],
+            'prediction': [1.1, 1.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['r2'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'r2'
+    assert result['value'].values[0] == 0.0  # Should return 0.0 when ss_tot == 0
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
+    )
+
+
+def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+            'target': [1.0, 2.0, 3.0],
+            'prediction': [1.1, 2.1, 2.9],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['rmse', 'mae'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 2
+    assert 'rmse' in result['metric'].values
+    assert 'mae' in result['metric'].values
+    assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
+    assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
+    assert all(data_size == 3 for data_size in result['data_size'].values)
+    assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
+    )
+
+
+def test_calculate_simple_metrics_unknown_metric_ignored(model_metrics_activity):
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['unknown_metric', 'rmse'],
+        'interval_minutes': 5,
+    }
+
+    result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
+
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'rmse'

tests/laborious/activities/test_opc.py

@@ -0,0 +1,720 @@
+from unittest.mock import ANY, MagicMock, call, patch
+
+import pytest
+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=MagicMock(),
+    )
+
+    assert opc.opc_servers == servers
+    assert opc.opc_repository == {}
+
+
+@patch('laborious.activities.opc.OpcRepository')
+@patch('laborious.activities.opc.OPC.send_notification')
+def test_init_opc(mock_send_notification, mock_opc_repository):
+    mock_logger = MagicMock()
+    mock_metrics_controller = MagicMock()
+    server1 = MagicMock(
+        connect=MagicMock(return_value=(True, {})), write_data=MagicMock(return_value=(True, {}))
+    )
+    server2 = MagicMock(
+        connect=MagicMock(return_value=(True, {})), write_data=MagicMock(return_value=(True, {}))
+    )
+    server3 = MagicMock(
+        connect=MagicMock(
+            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=MagicMock(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,
+    )
+    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.fixture
+@patch('laborious.activities.opc.OpcRepository')
+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 = MagicMock(return_value=(True, {}))
+    mock_opc_repository.return_value.connect = MagicMock(return_value=(True, {}))
+    opc = OPC(
+        opc_servers=servers,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=MagicMock(),
+    )
+    opc.init_opc()
+    opc.send_notification = MagicMock()
+    opc.emit_metric_sync = MagicMock()
+    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)
+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 = 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)
+
+
+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 = 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.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,
+    )
+
+
+def test_write_data_exception(opc):
+    opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
+
+    try:
+        opc.write_data(
+            server_id='server1',
+            tag='tag1',
+            data=50,
+            data_type='int',
+            tag_type='prediction',
+            metadata=metadata,
+        )
+
+    except Exception:
+        opc.send_notification.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
+
+
+def test_write_tags_from_config_prediction_success(opc):
+    opc.write_data = MagicMock(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 = 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'],
+    )
+
+
+def test_write_tags_from_config_confidence_success(opc):
+    opc.write_data = MagicMock(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 = 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'],
+    )
+
+
+def test_write_tags_from_config_write_failure(opc):
+    opc.write_data = MagicMock(return_value=(None, {}))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = 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
+
+
+def test_write_tags_from_config_session_bad(opc):
+    opc.write_data = MagicMock(
+        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 = 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
+
+
+def test_write_tags_from_config_reconnect_in_progress(opc):
+    opc.write_data = MagicMock(
+        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 = 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
+
+
+def test_manage_output_tags_success(opc):
+    opc._write_tags_from_config = MagicMock(
+        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 = 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.call_count == 2
+
+
+def test_manage_output_tags_failed(opc):
+    opc._write_tags_from_config = MagicMock(
+        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, _, _, _ = 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}
+
+
+def test_manage_output_tags_do_nothing(opc):
+    opc._write_tags_from_config = MagicMock()
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {'_invalid_key': {'tag1': {'data_type': 'float'}}}
+
+    output_data, opc_metrics, _, _, _ = 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()
+
+
+@patch('laborious.activities.opc.DataFrame')
+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 = MagicMock(
+        return_value=(True, {'tag1': 0.1, 'tag2': 0.2}, False, None, False)
+    )
+
+    opc.process_confidence = MagicMock(return_value={'data': 'data'})
+    output_data, opc_metrics = 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,
+    )
+
+
+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
+    opc.write_opc_data(input_data)
+
+    # Assert
+    opc.opc_repository['server1'].write_data.assert_not_called()
+
+
+def test_write_opc_data_no_validate_server(opc):
+    opc.validate_server = MagicMock(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
+    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
+
+
+def test_manage_output_tags_merges_error_flags(opc):
+    opc._write_tags_from_config = MagicMock(
+        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,
+    ) = 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]
+    )
+
+
+def test_validate_server(opc):
+    assert opc.validate_server('server1', metadata) is True
+    assert opc.validate_server('server2', metadata) is False
+
+
+def test_close(opc):
+    repo = opc.opc_repository['server1']
+    repo.disconnect = MagicMock(return_value=True)
+    opc.close()
+    repo.disconnect.assert_called_once()

tests/laborious/activities/test_storage.py

@@ -0,0 +1,324 @@
+import datetime
+import os
+from unittest.mock import ANY, MagicMock, patch
+
+from pytest import fixture, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+from sientia_do.temporal.activities.postgres_sync 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
+
+
+@fixture(autouse=True)
+def _patch_monitoring_shutdown():
+    """
+    Avoid running real async SientiaMonitoring.shutdown when Storage.close runs inside tests.
+    """
+    with patch.object(SientiaMonitoring, 'shutdown') as mock_shutdown:
+        yield mock_shutdown
+
+
+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=MagicMock(),
+    )
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___not_hasattr(mock_minio_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = MagicMock()
+    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 = MagicMock()
+    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=MagicMock(),
+    )
+    mock_minio_repository.assert_not_called()
+    assert storage.minio_repository is not None
+
+
+def test_close(storage, _patch_monitoring_shutdown):
+    storage.minio_repository = MagicMock()
+
+    storage.close()
+
+    assert storage.minio_repository is None
+    _patch_monitoring_shutdown.assert_called_once_with(storage)
+
+
+def test_close_when_minio_repository_already_none(storage, _patch_monitoring_shutdown):
+    """Closing without an initialized MinIO repository skips MinIO teardown."""
+    storage.minio_repository = None
+
+    storage.close()
+
+    assert storage.minio_repository is None
+    _patch_monitoring_shutdown.assert_called_once_with(storage)
+
+
+def test_load_query_with_minio_offload_no_rows(storage):
+    storage.load_custom_query = MagicMock(return_value=None)
+    storage_result = {'success': False}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=MagicMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = 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_called_once()
+
+
+def test_load_query_with_minio_offload_inline(storage):
+    storage.load_custom_query = MagicMock(return_value=[{'a': 1}])
+    storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=MagicMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = 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_called_once()
+
+
+def test_load_query_with_minio_offload_minio(storage):
+    storage.load_custom_query = MagicMock(return_value=[{'a': 1}])
+    storage_result = {'success': True, 'data': None, 'object_key': 'object-key'}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=MagicMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = 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_called_once()
+
+
+@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
+@patch('laborious.activities.storage.now')
+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 = MagicMock(
+        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 = MagicMock()
+    storage.send_notification = MagicMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+
+    result = 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'],
+    )
+
+
+def test_load_query_with_minio_offload_minio_not_initialized(storage):
+    storage.minio_repository = None
+
+    with raises(ValueError, match='Minio repository not initialized'):
+        storage.load_query_with_minio_offload({**metadata, 'query': 'SELECT 1', 'model_name': 'm'})
+
+
+def test_export_payload_to_postgres(storage):
+    payload = MagicMock()
+    payload.retrieve = MagicMock(return_value=MagicMock())
+    storage.export_data_to_postgres = MagicMock(return_value={'success': True})
+
+    result = storage.export_payload_to_postgres(
+        {**metadata, 'data': payload, 'schema': 'public', 'table': 't'}
+    )
+
+    payload.retrieve.assert_called_once_with(storage.minio_repository, metadata['metadata'])
+    storage.export_data_to_postgres.assert_called_once()
+    assert result == {'success': True}
+
+
+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'):
+        storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+
+@patch('laborious.activities.storage.now')
+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 = MagicMock(
+        return_value=['some/random/key-without-timestamp.parquet']
+    )
+    storage.minio_repository.delete_file = MagicMock()
+    storage.send_notification = MagicMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'test'
+    result = 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()
+
+
+@patch('laborious.activities.storage.now')
+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 = MagicMock(return_value=[old_key])
+    storage.minio_repository.delete_file = MagicMock(side_effect=Exception('delete error'))
+    storage.send_notification = MagicMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = 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'
+
+
+@patch('laborious.activities.storage.now')
+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 = MagicMock(side_effect=Exception('list error'))
+    storage.send_notification = MagicMock()
+    storage.error = MagicMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 0
+    storage.send_notification.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,278 @@
+from datetime import datetime
+from io import BytesIO
+from unittest.mock import MagicMock, patch
+
+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
+
+
+def test_retrieve_inline_dict_as_dataframe():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
+    minio = MagicMock()
+    out = payload.retrieve(minio, {'metadata': {}})
+    assert list(out.columns) == ['a']
+    minio.download_file.assert_not_called()
+
+
+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 = MagicMock()
+    minio.download_file = MagicMock(return_value=file_bytes)
+
+    out = payload.retrieve(minio, {'metadata': {}})
+
+    minio.download_file.assert_called_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
+
+
+def test_retrieve_empty_when_no_data():
+    payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
+    minio = MagicMock()
+    out = payload.retrieve(minio, {})
+    assert out.empty
+    minio.download_file.assert_not_called()
+
+
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+def test_from_dataframe_none(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = MagicMock()
+    result = 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
+
+
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+def test_from_dataframe_empty(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = MagicMock()
+    mock_df = MagicMock()
+    mock_df.__bool__ = MagicMock(return_value=True)
+    mock_df.empty = True
+    result = 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
+
+
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
+def test_from_dataframe_inline():
+    minio = MagicMock()
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = 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'
+
+
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
+def test_from_dataframe_inline_uses_provided_last_timestamp():
+    minio = MagicMock()
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = MinioDataFramePayload.from_dataframe(
+        dataframe=df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+        last_timestamp='2024-01-02',
+    )
+    assert result.last_timestamp == '2024-01-02'
+
+
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
+def test_from_dataframe_offloaded(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = MagicMock()
+    minio.upload_file = MagicMock(return_value={'minio_object_name': 'full/key.parquet'})
+    minio.bucket = 'test-bucket'
+
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = 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_called_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
+
+
+def test_debug_with_logger_calls_custom_debug():
+    logger = MagicMock()
+    MinioDataFramePayload._debug(logger, 'msg', {'a': 1})
+    logger.custom_debug.assert_called_once_with('msg', {'a': 1})

tests/laborious/utils/repository/test_opc_repository.py

@@ -0,0 +1,584 @@
+import concurrent.futures
+import json
+from datetime import datetime
+from unittest.mock import ANY, MagicMock, Mock, patch
+
+import pytest
+from asyncua.crypto import security_policies
+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=MagicMock(),
+    )
+    repository.disconnection_interval = 0.1
+    repository.send_notification = MagicMock()
+    repository.send_notification = MagicMock()
+    repository.emit_metric_sync = MagicMock()
+    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 = MagicMock()
+        aio = MagicMock()
+        client_instance.aio_obj = aio
+        aio.uaclient = MagicMock()
+        aio.uaclient.protocol = MagicMock(state='closed')
+        aio.session_timeout = 600_000
+        aio.secure_channel_timeout = 600_000
+        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
+
+
+def test_set_security(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    opc_repository.set_security()
+
+    mock_client.application_uri = 'urn:test:server'
+    mock_client.set_security.assert_called_once_with(
+        security_policies.SecurityPolicyBasic256,
+        '/path/to/cert.pem',
+        '/path/to/key.pem',
+        None,
+        '/path/to/server_cert.pem',
+    )
+    assert mock_client.aio_obj.secure_channel_timeout == 600_000
+    assert mock_client.aio_obj.session_timeout == 600_000
+
+
+def test_set_security_missing_certificates(opc_repository):
+    opc_repository.cert_path = None
+    opc_repository.private_key_path = None
+
+    try:
+        opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
+
+
+def test_set_security_missing_client(opc_repository):
+    opc_repository.client = None
+    try:
+        opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Client must be initialized before setting security'
+
+
+def test_connect_with_security(opc_repository, mock_client):
+    opc_repository._create_client = MagicMock()
+    opc_repository._open_session = MagicMock(return_value=(True, {}))
+    result = opc_repository.connect()
+
+    opc_repository._create_client.assert_called_once()
+    opc_repository._open_session.assert_called_once()
+    assert result == (True, {})
+
+
+def test_connect_without_security(opc_repository, mock_client):
+    opc_repository.cert_path = None
+    opc_repository._create_client = MagicMock()
+    opc_repository._open_session = MagicMock(return_value=(True, {}))
+    opc_repository.set_security = MagicMock()
+    result = 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, {})
+
+
+def test_connect_raises_when_session_already_open(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'open'
+    mock_client.aio_obj.uaclient.protocol = proto
+
+    with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
+        opc_repository.connect()
+
+
+def test_create_client_raises_when_client_exists(opc_repository, mock_client):
+    opc_repository.client = mock_client
+
+    with pytest.raises(OpcClientAlreadyExistsError, match='already exists'):
+        opc_repository._create_client()
+
+
+def test_open_session_success(opc_repository):
+    closed_proto = MagicMock()
+    closed_proto.state = 'closed'
+    opc_repository.client = MagicMock()
+    aio = MagicMock()
+    opc_repository.client.aio_obj = aio
+    aio.uaclient = MagicMock(protocol=closed_proto)
+    aio.session_timeout = 600_000
+    aio.secure_channel_timeout = 600_000
+
+    open_proto = MagicMock()
+    open_proto.state = 'open'
+    open_proto.authentication_token = 'tok'
+
+    def connect_side_effect():
+        aio.uaclient.protocol = open_proto
+
+    opc_repository.client.connect = MagicMock(side_effect=connect_side_effect)
+
+    result = opc_repository._open_session()
+
+    opc_repository.client.connect.assert_called_once()
+    assert result == (True, {})
+    assert opc_repository._session_ready.is_set()
+
+
+def test_open_session_raises_when_already_connected(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'open'
+    mock_client.aio_obj.uaclient.protocol = proto
+
+    with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
+        opc_repository._open_session()
+
+
+def test_open_session_fail(opc_repository):
+    opc_repository._disconnect_locked = MagicMock()
+    opc_repository.client = MagicMock()
+    aio = MagicMock()
+    opc_repository.client.aio_obj = aio
+    aio.uaclient = MagicMock(protocol=MagicMock(state='closed'))
+    opc_repository.client.connect = MagicMock(side_effect=Exception('Test error'))
+
+    is_connected, error_data = 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
+
+
+def test_open_session_raises_when_no_client(opc_repository):
+    opc_repository.client = None
+
+    with pytest.raises(OpcClientNotInitializedError, match='not initialized'):
+        opc_repository._open_session()
+
+
+def test_disconnection_fallback_success(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    mock_client.disconnect.return_value = True
+    result = opc_repository._disconnection_fallback()
+
+    mock_client.disconnect.assert_called_once()
+    assert result == []
+
+
+def test_disconnection_fallback_fail(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    mock_client.disconnect.side_effect = Exception('Test error')
+    result = 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
+
+
+def test_disconnect(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    opc_repository._disconnection_fallback = MagicMock(return_value=[])
+    opc_repository.disconnect()
+
+    opc_repository._disconnection_fallback.assert_called_once()
+    assert opc_repository.client is None
+    assert opc_repository._allow_reconnect is False
+
+
+def test_disconnect_no_client(opc_repository):
+    opc_repository.client = None
+    assert opc_repository.disconnect() is None
+
+
+def test_disconnect_error(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    opc_repository._disconnection_fallback = MagicMock(
+        return_value=[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}]
+    )
+    opc_repository.disconnect()
+
+    opc_repository._disconnection_fallback.assert_called_once()
+    opc_repository.send_notification.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
+
+
+def test_validate_connection_none_client(opc_repository):
+    opc_repository.client = None
+    response = opc_repository.validate_connection()
+    assert response == (False, opc_repository._not_connected_error())
+
+
+def test_validate_connection_session_not_open(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.aio_obj.uaclient.protocol = None
+
+    response = opc_repository.validate_connection()
+
+    assert response == (False, opc_repository._not_connected_error())
+    opc_repository.error.assert_called_once()
+
+
+def test_validate_connection_success(opc_repository):
+    opc_repository.client = MagicMock()
+    proto = MagicMock()
+    proto.state = 'open'
+    opc_repository.client.aio_obj.uaclient.protocol = proto
+
+    output = opc_repository.validate_connection()
+    assert output == (True, {})
+
+
+def test_write_data_validate_connection_do_nothing(opc_repository):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = MagicMock(get_node=MagicMock())
+    mock_node = MagicMock()
+    opc_repository.client.get_node.return_value = mock_node
+
+    result = 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})
+
+
+def test_write_data_validate_connection_failed(opc_repository):
+    opc_repository.validate_connection = MagicMock(return_value=(False, {}))
+    opc_repository.client = MagicMock()
+    opc_repository._start_reconnect = MagicMock()
+
+    is_success, error_data = 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_not_called()
+    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'
+
+
+def test_write_data_get_node_failed(opc_repository):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = MagicMock()
+    opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
+
+    is_success, error_data = 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
+
+
+def test_write_data_invalid_data_type(opc_repository, mock_client):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    mock_node = MagicMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+
+    is_success, error_data = 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
+
+
+def test_write_data(opc_repository, mock_client):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    mock_node = MagicMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+
+    result = 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})
+
+
+def test_write_data_write_value_failed(opc_repository, mock_client):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    mock_node = MagicMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = Exception('Test error')
+
+    is_success, error_data = 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
+
+
+def test_write_data_bad_session_id_invalid_schedules_reconnect(opc_repository, mock_client):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    opc_repository._start_reconnect = MagicMock()
+    mock_node = MagicMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = BadSessionIdInvalid()
+
+    is_success, error_data = 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'
+
+
+def test_write_data_reconnect_in_progress_immediate(opc_repository):
+    opc_repository._session_ready.clear()
+    opc_repository._reconnect_thread = MagicMock()
+    opc_repository._reconnect_thread.is_alive.return_value = True
+    opc_repository.validate_connection = MagicMock()
+
+    is_success, error_data = opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository.validate_connection.assert_not_called()
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'reconnect_in_progress'
+
+
+def test_start_reconnect_skips_within_interval(opc_repository):
+    opc_repository.last_reconnection_time = datetime.now()
+    opc_repository.reconnection_interval = 3600
+
+    opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
+
+    assert opc_repository._reconnect_thread is None
+
+
+def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.aio_obj.uaclient.protocol = MagicMock(state='closed')
+    opc_repository._start_reconnect = MagicMock()
+
+    is_success, error_data = 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'
+
+
+def test_write_data_protocol_closed_skips_reconnect_within_interval(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.aio_obj.uaclient.protocol = MagicMock(state='closed')
+    opc_repository.last_reconnection_time = datetime.now()
+    opc_repository.reconnection_interval = 3600
+
+    is_success, error_data = opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    assert opc_repository._reconnect_thread is None
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+
+
+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 = MagicMock(
+        return_value=(False, {'message': 'connect failed'})
+    )
+
+    opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+    if opc_repository._reconnect_thread is not None:
+        opc_repository._reconnect_thread.join(timeout=2)
+    assert opc_repository._reconnect_locked.call_count == 1
+    assert not opc_repository._reconnect_thread_in_progress()
+
+    opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+    if opc_repository._reconnect_thread is not None:
+        opc_repository._reconnect_thread.join(timeout=2)
+    assert opc_repository._reconnect_locked.call_count == 2
+
+
+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.aio_obj.uaclient.protocol = proto
+    opc_repository._disconnection_fallback = MagicMock(return_value=[])
+    opc_repository.disconnect()
+
+    is_success, error_data = opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    assert opc_repository._reconnect_thread is None
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+
+
+def test_parallel_bad_writes_single_reconnect_task(opc_repository, mock_client):
+    opc_repository.validate_connection = MagicMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    opc_repository.reconnection_interval = 0
+    opc_repository.last_reconnection_time = None
+    mock_node = MagicMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = BadSessionIdInvalid()
+    opc_repository._start_reconnect = MagicMock()
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [
+            executor.submit(
+                opc_repository.write_data,
+                node,
+                value,
+                'float',
+                metadata['metadata'],
+            )
+            for node, value in (('ns=2;s=TestNode', 1.0), ('ns=2;s=TestNode2', 2.0))
+        ]
+        results = [future.result() for future in futures]
+
+    assert 1 <= opc_repository._start_reconnect.call_count <= 2
+    assert 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)
+
+
+@patch('laborious.utils.repository.opc_repository.datetime')
+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 = MagicMock()
+    opc_repository._connect_locked = MagicMock(return_value=(True, {}))
+
+    result = 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,152 @@
+from os import environ
+
+from laborious.utils.connectors_config import (
+    build_minio_config,
+    build_mlflow_config,
+    build_opc_config,
+    build_plugin_store_config,
+)
+
+
+def test_build_mlflow_config_with_env_vars():
+    # Arrange
+    environ['MLFLOW_URL'] = 'http://test-host:8080'
+    environ['MLFLOW_USERNAME'] = 'test-user'
+    environ['MLFLOW_PASSWORD'] = 'test-pass'
+
+    # Act
+    config = build_mlflow_config()
+
+    # Assert
+    assert config['url'] == 'http://test-host:8080'
+    assert config['username'] == 'test-user'
+    assert config['password'] == 'test-pass'
+
+
+def test_build_mlflow_config_host_already_has_port():
+    environ['MLFLOW_URL'] = 'http://tracker.example.com:443'
+    environ['MLFLOW_USERNAME'] = 'u'
+    environ['MLFLOW_PASSWORD'] = 'p'
+
+    config = build_mlflow_config()
+
+    assert config['url'] == 'http://tracker.example.com:443'
+
+
+def test_build_mlflow_config_with_defaults():
+    # Arrange
+    # Clear any existing env vars
+    environ.pop('MLFLOW_URL', None)
+    environ.pop('MLFLOW_USERNAME', None)
+    environ.pop('MLFLOW_PASSWORD', None)
+
+    # Act
+    config = build_mlflow_config()
+
+    # Assert
+    assert config['url'] == 'http://localhost:5080'
+    assert config['username'] == 'aignosi'
+    assert config['password'] == 'aignosi'
+
+
+def test_build_plugin_store_config_defaults():
+    environ.pop('STORE_BASE_URL', None)
+    environ.pop('STORE_OWNER', None)
+    environ.pop('STORE_REPO', None)
+    environ.pop('STORE_BRANCH', None)
+    environ.pop('STORE_USERNAME', None)
+    environ.pop('STORE_PASSWORD', None)
+    environ.pop('STORE_CACHE_TTL_SECONDS', None)
+    environ.pop('PYPI_SERVER', None)
+    environ.pop('PYPI_USERNAME', None)
+    environ.pop('PYPI_PASSWORD', None)
+
+    cfg = build_plugin_store_config()
+    assert cfg['base_url'] == 'http://localhost:3000'
+    assert cfg['owner'] == 'sientia'
+    assert cfg['repo'] == 'model-library-store'
+    assert cfg['pypi_index_url'] == 'http://localhost:5000'
+
+
+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'
+    # Isolate from IDE/CI env (e.g. VS Code may export MINIO_SECURE=true).
+    environ['MINIO_SECURE'] = 'false'
+    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)
+    environ.pop('MINIO_SECURE', None)
+    environ.pop('MINIO_RETENTION_HOURS', 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/utils/test_dataframe_debug.py

@@ -0,0 +1,12 @@
+from pandas import DataFrame
+
+from laborious.utils.dataframe_debug import build_dataframe_debug_message
+
+
+def test_build_dataframe_debug_message_skips_large_dataframe():
+    df = DataFrame({'a': [1, 2, 3]})
+
+    msg = build_dataframe_debug_message('payload', df, max_rows=1)
+
+    assert 'skipped because dataframe has 3 rows' in msg
+    assert '(max: 1)' in msg

tests/laborious/worker/test_worker.py

@@ -0,0 +1,243 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from pytest import mark, raises
+
+from laborious.worker import worker
+
+
+def _build_fake_activities():
+    inst = MagicMock()
+    inst.init_opc = MagicMock()
+    inst.shutdown = MagicMock()
+    inst.load_query_with_minio_offload = MagicMock()
+    inst.retrain_model = MagicMock()
+    inst.update_production_model = MagicMock()
+    inst.format_retrain_report = MagicMock()
+    inst.export_data_to_postgres = MagicMock()
+    inst.load_custom_query = MagicMock()
+    inst.calculate_simple_metrics = MagicMock()
+    inst.get_reference_data = MagicMock()
+    inst.calculate_drift = MagicMock()
+    inst.request_predict = MagicMock()
+    inst.request_transform = MagicMock()
+    inst.input_gate = MagicMock()
+    inst.mlflow_response_gate = MagicMock()
+    inst.mlflow_content_gate = MagicMock()
+    inst.format_transformed_data = MagicMock()
+    inst.format_prediction = MagicMock()
+    inst.format_default_prediction = MagicMock()
+    inst.write_opc_data = MagicMock()
+    inst.cleanup_minio_objects_expired = MagicMock()
+    inst.repeat_last_prediction = MagicMock()
+    inst.write_metrics = MagicMock()
+    inst.write_pi_web_api_data = MagicMock()
+    return inst
+
+
+def _build_fake_worker(async_result=None, async_error: Exception | None = None):
+    w = MagicMock()
+
+    async def _run():
+        if async_error is not None:
+            raise async_error
+        return async_result
+
+    w.run = MagicMock(side_effect=_run)
+    return w
+
+
+@patch('laborious.worker.worker.start_http_server')
+def test_start_prometheus_server_success(mock_start_http):
+    with patch.object(worker.metrics.APP_UP, 'labels') as labels:
+        gauge = MagicMock()
+        labels.return_value = gauge
+        with patch('laborious.worker.worker.os.getenv', return_value='9090'):
+            worker.start_prometheus_server()
+    mock_start_http.assert_called_once_with(9090)
+    gauge.set.assert_called_once_with(1)
+
+
+@patch('laborious.worker.worker.start_http_server', side_effect=RuntimeError('nope'))
+def test_start_prometheus_server_error_exits(_mock_start_http):
+    with patch('laborious.worker.worker.os._exit', side_effect=SystemExit(1)) as m_exit:
+        with raises(SystemExit):
+            worker.start_prometheus_server()
+    m_exit.assert_called_once_with(1)
+
+
+@mark.asyncio
+async def test_main_missing_runtime_exits_fast(monkeypatch):
+    monkeypatch.setenv('RUNTIME', '')
+    with (
+        patch('laborious.worker.worker.start_prometheus_server'),
+        patch('laborious.worker.worker.get_logger') as m_logger,
+        patch('laborious.worker.worker.NotificationHandler'),
+        patch('laborious.worker.worker.MetricsController'),
+        patch.object(worker.metrics.APP_UP, 'labels') as labels,
+        patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
+    ):
+        labels.return_value = MagicMock()
+        with raises(SystemExit):
+            await worker.main()
+    assert m_logger.return_value.custom_critical.called
+
+
+@mark.asyncio
+async def test_main_plugin_install_failure(monkeypatch):
+    monkeypatch.setenv('RUNTIME', 'single')
+    fake_activities = _build_fake_activities()
+    fake_plugin = MagicMock()
+    fake_plugin.install_runtime = AsyncMock(side_effect=RuntimeError('install failed'))
+    with (
+        patch('laborious.worker.worker.start_prometheus_server'),
+        patch('laborious.worker.worker.get_logger'),
+        patch(
+            'laborious.worker.worker.build_mongodb_config',
+            return_value={'connection_string': 'cs', 'database_name': 'db'},
+        ),
+        patch('laborious.worker.worker.NotificationHandler'),
+        patch('laborious.worker.worker.MetricsController'),
+        patch(
+            'laborious.worker.worker.build_plugin_store_config',
+            return_value={
+                'base_url': '',
+                'owner': '',
+                'repo': '',
+                'username': None,
+                'password': None,
+                'branch': None,
+                'cache_ttl_seconds': None,
+                'pypi_index_url': '',
+                'pypi_username': None,
+                'pypi_password': None,
+            },
+        ),
+        patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
+        patch('laborious.worker.worker.Activities', return_value=fake_activities),
+        patch.object(worker.metrics.APP_UP, 'labels') as labels,
+        patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
+    ):
+        labels.return_value = MagicMock()
+        with raises(SystemExit):
+            await worker.main()
+
+
+@mark.asyncio
+async def test_main_success_exit_zero(monkeypatch):
+    monkeypatch.setenv('RUNTIME', 'single')
+    fake_activities = _build_fake_activities()
+    fake_plugin = MagicMock()
+    fake_plugin.install_runtime = AsyncMock(return_value=None)
+    fake_workers = [_build_fake_worker() for _ in range(4)]
+
+    with (
+        patch('laborious.worker.worker.start_prometheus_server'),
+        patch('laborious.worker.worker.get_logger'),
+        patch(
+            'laborious.worker.worker.build_mongodb_config',
+            return_value={'connection_string': 'cs', 'database_name': 'db'},
+        ),
+        patch('laborious.worker.worker.NotificationHandler') as m_notif_cls,
+        patch('laborious.worker.worker.MetricsController'),
+        patch(
+            'laborious.worker.worker.build_plugin_store_config',
+            return_value={
+                'base_url': '',
+                'owner': '',
+                'repo': '',
+                'username': None,
+                'password': None,
+                'branch': None,
+                'cache_ttl_seconds': None,
+                'pypi_index_url': '',
+                'pypi_username': None,
+                'pypi_password': None,
+            },
+        ),
+        patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
+        patch('laborious.worker.worker.Activities', return_value=fake_activities),
+        patch('laborious.worker.worker.build_postgres_config', return_value={}),
+        patch('laborious.worker.worker.build_minio_config', return_value={}),
+        patch('laborious.worker.worker.build_opc_config', return_value={}),
+        patch('laborious.worker.worker.build_api_config', return_value={}),
+        patch('laborious.worker.worker.PrometheusConfig', return_value=MagicMock()),
+        patch('laborious.worker.worker.TelemetryConfig', return_value=MagicMock()),
+        patch('laborious.worker.worker.Runtime', return_value=MagicMock()),
+        patch(
+            'laborious.worker.worker.client.Client.connect', new=AsyncMock(return_value=MagicMock())
+        ),
+        patch('laborious.worker.worker.prepare_worker', side_effect=fake_workers) as m_prepare,
+        patch.object(worker.metrics.APP_UP, 'labels') as labels,
+        patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(0)),
+    ):
+        labels.return_value = MagicMock()
+        with raises(SystemExit):
+            await worker.main()
+
+    notif = m_notif_cls.return_value
+    notif.shutdown.assert_called_once()
+    fake_activities.shutdown.assert_called_once()
+    assert m_prepare.call_count == 4
+    prepare_calls = m_prepare.call_args_list
+    assert prepare_calls[0].kwargs['runtime'] == 'single'
+    assert prepare_calls[3].kwargs['runtime'] == 'single'
+
+
+@mark.asyncio
+async def test_main_worker_gather_error_exits_one(monkeypatch):
+    monkeypatch.setenv('RUNTIME', 'single')
+    fake_activities = _build_fake_activities()
+    fake_plugin = MagicMock()
+    fake_plugin.install_runtime = AsyncMock(return_value=None)
+    fake_workers = [
+        _build_fake_worker(async_error=RuntimeError('boom')),
+        _build_fake_worker(),
+        _build_fake_worker(),
+        _build_fake_worker(),
+    ]
+
+    with (
+        patch('laborious.worker.worker.start_prometheus_server'),
+        patch('laborious.worker.worker.get_logger') as m_logger,
+        patch(
+            'laborious.worker.worker.build_mongodb_config',
+            return_value={'connection_string': 'cs', 'database_name': 'db'},
+        ),
+        patch('laborious.worker.worker.NotificationHandler'),
+        patch('laborious.worker.worker.MetricsController'),
+        patch(
+            'laborious.worker.worker.build_plugin_store_config',
+            return_value={
+                'base_url': '',
+                'owner': '',
+                'repo': '',
+                'username': None,
+                'password': None,
+                'branch': None,
+                'cache_ttl_seconds': None,
+                'pypi_index_url': '',
+                'pypi_username': None,
+                'pypi_password': None,
+            },
+        ),
+        patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
+        patch('laborious.worker.worker.Activities', return_value=fake_activities),
+        patch('laborious.worker.worker.build_postgres_config', return_value={}),
+        patch('laborious.worker.worker.build_minio_config', return_value={}),
+        patch('laborious.worker.worker.build_opc_config', return_value={}),
+        patch('laborious.worker.worker.build_api_config', return_value={}),
+        patch('laborious.worker.worker.PrometheusConfig', return_value=MagicMock()),
+        patch('laborious.worker.worker.TelemetryConfig', return_value=MagicMock()),
+        patch('laborious.worker.worker.Runtime', return_value=MagicMock()),
+        patch(
+            'laborious.worker.worker.client.Client.connect', new=AsyncMock(return_value=MagicMock())
+        ),
+        patch('laborious.worker.worker.prepare_worker', side_effect=fake_workers),
+        patch.object(worker.metrics.APP_UP, 'labels') as labels,
+        patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
+    ):
+        labels.return_value = MagicMock()
+        with raises(SystemExit):
+            await worker.main()
+
+    assert m_logger.return_value.custom_error.called

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,866 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+from pytest import fixture, mark, raises
+
+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,
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
+async def test_run_always_cleans_up_on_pipeline_exception(workflow_mock, prediction_process):
+    input_data = {
+        'metadata': metadata,
+        'data': {'last_timestamp': '2024-01-01'},
+        'model_id': 1,
+        'model_name': 'm',
+        'model_config': {},
+        'save_transform': False,
+    }
+    prediction_process._run_prediction_pipeline = AsyncMock(side_effect=RuntimeError('boom'))
+
+    with raises(RuntimeError):
+        await prediction_process.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.cleanup_minio_objects_expired,
+        {**metadata, 'data': input_data['data']},
+        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,324 @@
+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',
+            'retention_minutes': 0,
+        },
+    }
+
+    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',
+            'retention_minutes': 0,
+        },
+    }
+
+    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',
+            'retention_minutes': 0,
+        },
+    }
+
+    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,
+    )