From 07dc61211613f46ce660e4a4af3e9392612d2695 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 16:52:33 -0300 Subject: [PATCH] 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 | 53 +++++-- laborious/activities/gates.py | 67 +++++++- laborious/activities/mlflow.py | 18 ++- .../format_and_export_prediction.py | 28 +++- tests/laborious/activities/test_gates.py | 146 +++++++++++++++++ tests/laborious/activities/test_mlflow.py | 89 +++++++++++ .../test_format_and_export_prediction.py | 150 ++++++++++++++++++ 7 files changed, 531 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 0143cbd..a34b939 100644 --- a/README.md +++ b/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 diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 13072cf..f0f8474 100644 --- a/laborious/activities/gates.py +++ b/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) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 6eadf60..3970bb7 100644 --- a/laborious/activities/mlflow.py +++ b/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'] diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 1851c61..e50c2cf 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/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, { diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 8399b43..e33ea4f 100644 --- a/tests/laborious/activities/test_gates.py +++ b/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' diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 93d71ee..dd7a0c5 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/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'], + ) diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index ecd3bc2..f5619f3 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/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',