SIENTIAPDE-1273

Enhance data handling and export processes in Laborious workflows

- Updated `gates.py` to improve data quality validation, filtering, and formatting operations, including enhanced metrics recording.
- Refined `mlflow.py` to better manage model transformations and reference data retrieval from MLflow Model Registry.
- Enhanced `format_and_export_prediction.py` to support separate export of transformed data, improving flexibility in data handling.
- Added comprehensive test coverage for new functionalities, including transformed data formatting and retrain report generation.
- Improved documentation in `README.md` to reflect changes in activities and workflows, ensuring clarity on data processing and export paths.

README.md

@@ -130,8 +130,14 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
 - `minimal_retrain.py`: Automated model retraining and production update
 
 #### **Activities (`laborious/activities/`)**
-- `gates.py`: Data quality validation and filtering
-- `mlflow.py`: Transform and predict operations
+- `gates.py`: Data quality validation, filtering, and data formatting operations
+  - Input/response/content gates for quality validation
+  - Prediction and transformed data formatting
+  - Retrain report formatting and metrics recording
+- `mlflow.py`: Transform, predict, and model management operations
+  - MLFlow model transformation and prediction
+  - Model retraining and production updates
+  - Reference data retrieval from MLflow Model Registry
 - `opc.py`: OPC UA export to industrial systems (optional)
 - `activities.py`: Aggregates activity interfaces
 
@@ -146,7 +152,9 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
 #### **1. Batch Prediction Pipeline**
 ```
 Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform →
-MLFlow Prediction → Response Validation → Export (PostgreSQL [+ OPC])
+MLFlow Prediction → Response Validation → Format & Export
+  ├─→ Predictions → PostgreSQL [+ OPC]
+  └─→ Transformed Data → PostgreSQL (optional)
 ```
 
 #### **2. Model Retraining Pipeline**
@@ -318,27 +326,37 @@ The **FormatAndExportPrediction** workflow handles prediction data formatting an
 #### Execution Flow
 1. **Path Decision**: Determines formatting path based on configuration
 2. **Data Formatting**: Formats prediction data for specific output requirements
-3. **PostgreSQL Export**: Writes formatted predictions to database
-4. **OPC Export**: Writes predictions to OPC servers
-5. **Metrics Recording**: Records export performance and success metrics
+3. **Transformed Data Processing**: Optionally formats and exports transformed data separately
+4. **PostgreSQL Export**: Writes formatted predictions to database
+5. **OPC Export**: Writes predictions to OPC servers
+6. **Metrics Recording**: Records export performance and success metrics
 
 #### Key Features
 - **Flexible Formatting**: Configurable output formats for different destinations
 - **Multi-Destination Export**: PostgreSQL and OPC server integration
+- **Transformed Data Export**: Optional separate export of MLFlow transformed data
 - **Performance Monitoring**: Comprehensive metrics for export operations
 - **Error Handling**: Robust error handling with notification integration
 
 #### Architecture Diagram
 ```mermaid
 flowchart LR
-    A[1. format_prediction/format_default_prediction] --> B[2. write_opc_data] --> C[3. export_data_to_postgres] --> D[4. write_metrics]
+    A[1. format_prediction/format_default_prediction] --> B[2. format_transformed_data] --> C[3. write_opc_data] --> D[4. export_data_to_postgres] --> E[5. write_metrics]
     
     A -.-> Format[Data Formatting]
-    B -.-> OPC[OPC Servers]
-    C -.-> PostgreSQL[(PostgreSQL)]
-    D -.-> Prometheus[Prometheus]
+    B -.-> Transform[Transformed Data]
+    C -.-> OPC[OPC Servers]
+    D -.-> PostgreSQL[(PostgreSQL)]
+    E -.-> Prometheus[Prometheus]
 ```
 
+#### Transformed Data Export
+When `transformed_data` is provided in the input, the workflow will:
+- Format the transformed data using `format_transformed_data` activity
+- Export it to a separate table (`transform_table_name`) asynchronously
+- Wait for both prediction and transformed data exports to complete
+- This enables separate tracking of model transformations for analysis and debugging
+
 ### 4. Minimal Retrain Workflow (`minimal_retrain.py`)
 
 The **MinimalRetrain** workflow handles automated model retraining and production model updates.
@@ -579,11 +597,24 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac
 ```
 tests/
 ├── activities/           # Activity implementation tests
-├── workflow/            # Workflow orchestration tests
+│   ├── test_gates.py    # Data quality gates and formatting tests
+│   ├── test_mlflow.py   # MLFlow operations and reference data tests
+│   └── ...              # Other activity tests
+├── workflows/           # Workflow orchestration tests
+│   └── subworkflows/    # Sub-workflow tests
+│       └── test_format_and_export_prediction.py  # Export workflow tests
 ├── utils/               # Utility function tests
 └── integration/         # End-to-end workflow tests
 ```
 
+### Test Coverage
+The test suite provides comprehensive coverage for:
+- **Data Quality Gates**: Input, response, and content validation filters
+- **Data Formatting**: Prediction, transformed data, and retrain report formatting
+- **MLFlow Operations**: Transform, predict, retrain, and reference data retrieval
+- **Workflow Orchestration**: Complete workflow execution paths and error handling
+- **Metrics Recording**: Performance monitoring and OPC export metrics
+
 ### Test Execution
 ```bash
 # Install test dependencies

laborious/activities/gates.py

@@ -405,7 +405,32 @@ class Gates(SientiaMonitoring):
     @activity.defn(name='format_transformed_data')
     async def format_transformed_data(self, input_data: dict[str, Any]) -> dict:
         """
-        Format transformed data according to configured storage policies.
+        Format transformed data for storage and export operations.
+
+        This method formats transformed data from MLFlow model transformations
+        into a standardized format suitable for database storage. It converts
+        wide-format data (columns as variables) into long-format (melted)
+        with proper timestamp handling and model identification.
+
+        The formatting process includes:
+        1. Converting input data dictionary to DataFrame
+        2. Extracting timestamps from DataFrame index
+        3. Resetting index to create sequential row numbers
+        4. Melting data from wide format to long format (variable-value pairs)
+        5. Adding model_id for data lineage tracking
+
+        Args:
+            input_data (dict): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict)
+                - model_id (str): Unique identifier for the ML model
+
+        Returns:
+            dict: Formatted data dictionary with keys:
+                - timestamp (dict): Timestamp values indexed by row number
+                - variable (dict): Variable names indexed by row number
+                - value (dict): Variable values indexed by row number
+                - model_id (dict): Model identifiers indexed by row number
         """
         metadata = input_data['metadata']
 
@@ -544,7 +569,45 @@ class Gates(SientiaMonitoring):
     @activity.defn(name='format_retrain_report')
     async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
         """
-        Format retrain report data according to configured storage policies.
+        Format retrain report data for storage and audit trail maintenance.
+
+        This method formats model retraining operation results into a standardized
+        report format suitable for database storage and operational monitoring.
+        It captures retraining status, timestamps, and model version information
+        for comprehensive audit trails and operational visibility.
+
+        The formatting process includes:
+        1. Extracting retraining experiment response data
+        2. Capturing model update report information (version, MLflow IDs)
+        3. Formatting timestamps and status information
+        4. Conditionally including version information for successful retrains
+
+        Args:
+            input_data (dict): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - experiment_response (dict): Retraining experiment response containing:
+                    - success (bool): Retraining operation success status
+                    - timestamp (str): Timestamp of the retraining operation
+                    - message (str): Status message or error description
+                - update_report (dict): Model update report containing:
+                    - version (str): New model version identifier
+                    - mlflow_run_id (str): MLflow run identifier
+                    - mlflow_experiment_id (str): MLflow experiment identifier
+                - model_id (str): Unique identifier for the ML model
+                - model_name (str): Name of the ML model
+
+        Returns:
+            dict: Formatted retrain report dictionary with keys:
+                - model_id (dict): Model identifiers indexed by row number
+                - model_name (dict): Model names indexed by row number
+                - timestamp (dict): Retraining timestamps indexed by row number
+                - status (dict): Retraining status messages indexed by row number
+                - version (dict, optional): Model versions indexed by row number
+                    Only included if experiment_response['success'] is True
+                - mlflow_run_id (dict, optional): MLflow run IDs indexed by row number
+                    Only included if experiment_response['success'] is True
+                - mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number
+                    Only included if experiment_response['success'] is True
         """
         metadata = input_data['metadata']
         self.info('Formatting retrain report...', metadata)

laborious/activities/mlflow.py

@@ -424,14 +424,30 @@ class MLFlow(SientiaMonitoring):
         """
         Get reference data from the MLflow Model Registry.
 
+        This method retrieves evaluation reference data stored as artifacts in the
+        MLflow Model Registry. The reference data is typically used for model
+        drift detection, performance comparison, and quality validation. The method
+        loads the data from a CSV artifact file and formats timestamps for
+        consistent processing.
+
+        The method handles:
+        1. Loading evaluation data artifact from MLflow Model Registry
+        2. Timestamp parsing and formatting for consistency
+        3. Data conversion to dictionary format for workflow consumption
+        4. Graceful handling of missing reference data
+
         Args:
             input_data (dict): Input data containing:
                 - metadata (dict): Workflow execution metadata
                 - model_name (str): Name of the MLFlow model to get reference data from
 
         Returns:
-            list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry.
+            list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
+                as a list of dictionaries. Returns None if reference data is not found
+                or if the artifact does not exist.
 
+        Raises:
+            Exception: If artifact loading fails or encounters errors during processing
         """
 
         metadata = input_data['metadata']

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -50,21 +50,36 @@ class FormatAndExportPrediction:
         Args:
             input_data: Complete configuration for the export workflow
                 Required keys:
+                    - metadata (dict): Workflow execution metadata
                     - path_flag (str | None): Decision path flag for formatting strategy
+                        - None: Normal prediction path with full formatting
+                        - Any other value: Default prediction path for error conditions
                     - data (dict[str, Any]): Prediction data to format and export
                     - prediction_confidence (float): Confidence score for the prediction
                     - timestamp (str): ISO-formatted timestamp for the prediction
                     - model_id (int): Unique identifier for the ML model
                     - model_name (str): Name of the ML model
-                    - model_retention (str): Model retention policy configuration
-                    - comment (str): Operational comment or error description
                     - schema (str): Database schema for data storage
                     - table_name (str): Target table for data persistence
                     - opc_output_config (dict[str, Any]): OPC server export configuration
-                    - prediction_store_policy (str, optional): Data retention policy
+                Optional keys:
+                    - transformed_data (dict[str, Any]): Transformed data to export separately
+                        Only processed when path_flag is None
+                    - transform_table_name (str): Target table for transformed data export
+                        Required if transformed_data is provided
+                    - prediction_store_policy (str): Data retention policy (e.g., 'lts:1', 'erl:2')
+                        Required when path_flag is None
+                    - comment (str): Operational comment or error description
+                        Required when path_flag is not None
 
         Returns:
-            bool: True if the workflow completes successfully, False otherwise
+            None: The workflow completes successfully when all export operations finish
+
+        Note:
+            When transformed_data is provided and path_flag is None, the workflow will:
+            1. Format the transformed data using format_transformed_data
+            2. Export it to a separate table (transform_table_name) asynchronously
+            3. Wait for both prediction and transformed data exports to complete
         """
         metadata = input_data['metadata']
         path_flag = input_data['path_flag']
@@ -73,7 +88,7 @@ class FormatAndExportPrediction:
         prediction_confidence = input_data['prediction_confidence']
 
         if path_flag is None:
-            # proceed with formatting and exporting
+            # Normal prediction path: format prediction data with full metadata
             prediction = await workflow.execute_local_activity_method(
                 Activities.format_prediction,
                 {
@@ -88,6 +103,7 @@ class FormatAndExportPrediction:
                 start_to_close_timeout=timedelta(seconds=60),
             )
 
+            # Optionally format and export transformed data to separate table
             if transformed_data is not None:
                 transformed = await workflow.execute_local_activity_method(
                     Activities.format_transformed_data,
@@ -120,7 +136,7 @@ class FormatAndExportPrediction:
                 write_transformed_handler = None
 
         else:
-            # create default prediction
+            # Error path: create default prediction with error indicators
             prediction = await workflow.execute_local_activity_method(
                 Activities.format_default_prediction,
                 {

tests/laborious/activities/test_gates.py

@@ -546,6 +546,80 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
         raise AssertionError('Expected ValueError')
 
 
+@mark.asyncio
+async def test_format_transformed_data_single_row(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {
+            'var1': {'2023-05-26 11:12:27': 1.0},
+            'var2': {'2023-05-26 11:12:27': 2.0},
+        },
+        'model_id': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_transformed_data(input_data)
+
+    # Assert
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27', 1: '2023-05-26 11:12:27'}
+    assert result['variable'] == {0: 'var1', 1: 'var2'}
+    assert result['value'] == {0: 1.0, 1: 2.0}
+    assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
+    gates_activity.info.assert_called()
+
+
+@mark.asyncio
+async def test_format_transformed_data_multiple_rows(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {
+            'var1': {
+                '2023-05-26 11:12:27': 1.0,
+                '2023-05-26 11:12:28': 2.0,
+            },
+            'var2': {
+                '2023-05-26 11:12:27': 3.0,
+                '2023-05-26 11:12:28': 4.0,
+            },
+        },
+        'model_id': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_transformed_data(input_data)
+
+    # Assert
+    assert len(result['timestamp']) == 4
+    assert len(result['variable']) == 4
+    assert len(result['value']) == 4
+    assert len(result['model_id']) == 4
+    assert all(v == 'test_model' for v in result['model_id'].values())
+    assert set(result['variable'].values()) == {'var1', 'var2'}
+    gates_activity.info.assert_called()
+
+
+@mark.asyncio
+async def test_format_transformed_data_empty_data(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {},
+        'model_id': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_transformed_data(input_data)
+
+    # Assert
+    assert result['timestamp'] == {}
+    assert result['variable'] == {}
+    assert result['value'] == {}
+    assert result['model_id'] == {}
+    gates_activity.info.assert_called()
+
+
 @mark.asyncio
 async def test_format_default_prediction(gates_activity):
     # Arrange
@@ -603,6 +677,40 @@ async def test_format_retrain_report(gates_activity):
     assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'}
 
 
+@mark.asyncio
+async def test_format_retrain_report_failure(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'experiment_response': {
+            'success': False,
+            'timestamp': '2023-05-26 11:12:27',
+            'message': 'failure',
+        },
+        'update_report': {
+            'version': '1.0.0',
+            'mlflow_run_id': 'test_mlflow_run_id',
+            'mlflow_experiment_id': 'test_mlflow_experiment_id',
+        },
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_retrain_report(input_data)
+
+    # Assert
+    assert result['model_id'] == {0: 'test_model'}
+    assert result['model_name'] == {0: 'test_model'}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
+    assert result['status'] == {0: 'failure'}
+    assert 'version' not in result
+    assert 'mlflow_run_id' not in result
+    assert 'mlflow_experiment_id' not in result
+    gates_activity.info.assert_called()
+    gates_activity.debug.assert_called()
+
+
 @mark.asyncio
 async def test_get_last_timestamp_with_data(gates_activity):
     # Arrange
@@ -742,3 +850,41 @@ async def test_write_metrics(mock_metrics, gates_activity):
             ),
         ]
     )
+
+
+@mark.asyncio
+@patch('laborious.activities.gates.metrics')
+async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_activity):
+    """Test write_metrics method with None response_time in opc_metrics."""
+    input_data = {
+        **metadata,
+        'prediction': {
+            'prediction': [1],
+            'prediction_confidence': [0.9],
+            'response_time': [0.1],
+        },
+        'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': None}},
+    }
+    await gates_activity.write_metrics(input_data)
+
+    # Verify that metrics for tag1 are emitted
+    gates_activity.emit_metric.assert_any_call(
+        metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
+        method='observe',
+        tags={
+            'pod_id': gates_activity.pod_id,
+            'model_name': metadata['metadata']['model_name'],
+            'workflow_name': metadata['metadata']['workflow_name'],
+            'opc_server_id': 'server1',
+            'tag': 'tag1',
+        },
+        value=0.1,
+    )
+
+    # Verify that metrics for tag2 (with None response_time) are NOT emitted
+    calls = [
+        c
+        for c in gates_activity.emit_metric.call_args_list
+        if len(c[1].get('tags', {})) > 0 and c[1]['tags'].get('tag') == 'tag2'
+    ]
+    assert len(calls) == 0, 'Metrics should not be emitted for None response_time'

tests/laborious/activities/test_mlflow.py

@@ -76,6 +76,11 @@ def mlflow(mock_minio_repository, mock_mlflow_repository):
     mlflow.send_notification = MagicMock()
     mlflow.emit_metric = AsyncMock()
     mlflow.send_notification_async = AsyncMock()
+    mlflow.error = MagicMock()
+    mlflow.debug = MagicMock()
+    mlflow.info = MagicMock()
+    mlflow.warning = MagicMock()
+    mlflow.critical = MagicMock()
 
     return mlflow
 
@@ -488,3 +493,87 @@ async def test_update_production_model_error(mlflow):
         )
     else:
         raise AssertionError('No exception raised')
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_get_reference_data_success(mock_to_datetime, mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    # Mock reference data DataFrame
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [
+        {'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
+        {'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
+    ]
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
+
+    # Act
+    result = await mlflow.get_reference_data(input_data)
+
+    # Assert
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )
+    mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
+
+    mock_reference_data.to_dict.assert_called_once_with(orient='records')
+    assert result == mock_reference_data.to_dict.return_value
+
+
+@mark.asyncio
+async def test_get_reference_data_not_found(mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
+
+    # Act
+    result = await mlflow.get_reference_data(input_data)
+
+    # Assert
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )
+    mlflow.warning.assert_called_once_with(
+        'Reference data not found for model test_model', metadata['metadata']
+    )
+    assert result is None
+
+
+@mark.asyncio
+async def test_get_reference_data_exception(mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
+        'Error loading artifact'
+    )
+
+    # Act & Assert
+    with raises(Exception) as e:
+        await mlflow.get_reference_data(input_data)
+
+    assert str(e.value) == 'Error loading artifact'
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -125,6 +125,156 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
     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,
+        '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,
+                {
+                    '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'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.format_transformed_data,
+                {
+                    'data': input_data['transformed_data'],
+                    'model_id': input_data['model_id'],
+                    **metadata,
+                },
+                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_data_to_postgres,
+        {
+            'schema': input_data['schema'],
+            'table_name': input_data['transform_table_name'],
+            'data': transformed_data,
+            'timestamp_conversion': {
+                'column': 'timestamp',
+                'format': DATETIME_FORMAT_WITH_TZ,
+            },
+            **metadata,
+        },
+        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,
+                {
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    **metadata,
+                },
+                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',