SIENTIAPDE-1478

Enhance end-to-end tests for PredictionsBatch workflow scenarios

- Updated scenario descriptions and assertions for error handling in prediction exports.
- Introduced new test cases for handling exports with only OPC or PI Web API.
- Refactored existing tests to improve clarity and maintainability, including dynamic data insertion.
- Adjusted workflow input configurations to better reflect expected behaviors for various error scenarios.

e2e/scenarios.md

@@ -326,48 +326,6 @@ The `predictions_batch` workflow:
 
 ---
 
-### 2.4 Error Scenarios
-
-#### Scenario 2.4.1: MLFlow Transform API Error
-**Description**: MLFlow transform request fails
-
-**Input**:
-- Valid input data
-- MLFlow service unavailable or returns error
-
-**Expected Behavior**:
-- `request_transform` raises exception
-- Notification sent with MLFlow error details
-- Workflow fails after retry attempts
-
-**Assertions**:
-- Exception raised from transform activity
-- Error notification sent
-- Workflow fails
-- Export NOT called
-
----
-
-#### Scenario 2.4.2: MLFlow Predict API Error
-**Description**: MLFlow predict request fails
-
-**Input**:
-- Valid input and transform data
-- MLFlow predict service unavailable
-
-**Expected Behavior**:
-- `request_predict` raises exception
-- Notification sent
-- Workflow fails after retries
-
-**Assertions**:
-- Transform succeeded
-- Predict raised exception
-- Error notification sent
-- Workflow fails
-
----
-
 ## 3. Format and Export Prediction - Child Workflow Scenarios
 
 ### 3.1 Success Scenarios
@@ -376,7 +334,7 @@ The `predictions_batch` workflow:
 **Description**: Error prediction path creates default prediction
 
 **Input**:
-- `path_flag: 'STOP'` or other non-None value
+- `path_flag: 'ERROR'` or other non-None value (not STOP/CONTINUE/REPEAT)
 - `comment` provided with error details
 
 **Expected Behavior**:
@@ -396,7 +354,53 @@ The `predictions_batch` workflow:
 
 ---
 
-#### Scenario 3.1.2: Export Without Optional Outputs
+#### Scenario 3.1.2: Export with OPC only
+**Description**: Export to PostgreSQL and OPC server only (no PI Web API)
+
+**Input**:
+- `path_flag: None`
+- `opc_output_config` configured with valid OPC settings
+- `pi_web_api_output_config: None` or `{}`
+
+**Expected Behavior**:
+- Normal formatting
+- PostgreSQL export executed
+- OPC export executed
+- PI Web API activity skipped
+- Metrics written with OPC metrics
+
+**Assertions**:
+- PI Web API activity NOT called
+- OPC activity called
+- PostgreSQL export called
+- Metrics written with `opc_metrics` populated
+
+---
+
+#### Scenario 3.1.3: Export with PI Web API only
+**Description**: Export to PostgreSQL and PI Web API only (no OPC)
+
+**Input**:
+- `path_flag: None`
+- `pi_web_api_output_config` configured with valid PI Web API settings
+- `opc_output_config: None` or `{}`
+
+**Expected Behavior**:
+- Normal formatting
+- PostgreSQL export executed
+- PI Web API export executed
+- OPC activity skipped
+- Metrics written without OPC metrics
+
+**Assertions**:
+- OPC activity NOT called
+- PI Web API activity called
+- PostgreSQL export called
+- Metrics written with empty `opc_metrics`
+
+---
+
+#### Scenario 3.1.4: Export Without Optional Outputs
 **Description**: Export only to PostgreSQL (no OPC or PI Web API)
 
 **Input**:
@@ -418,12 +422,14 @@ The `predictions_batch` workflow:
 
 ---
 
-#### Scenario 3.1.3: Export Without Transformed Data
+#### Scenario 3.1.5: Export Without Transformed Data
 **Description**: Only prediction exported, no transform table
 
 **Input**:
 - `path_flag: None`
-- `transformed_data: None`
+- `transformed_data: None` or `save_transform: False`
+- `opc_output_config: None` or `{}`
+- `pi_web_api_output_config: None` or `{}`
 
 **Expected Behavior**:
 - Only prediction formatted and exported

e2e/test_predictions_batch_format_export.py

@@ -4,6 +4,7 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
 
 import asyncio
 from datetime import datetime
+from unittest.mock import patch
 
 import pandas as pd
 import pytest
@@ -14,6 +15,78 @@ from temporalio.worker import Worker
 from laborious.activities.activities import Activities
 from laborious.workflows.predictions_batch import PredictionsBatch
 
+base_input_data = {
+    'schedule_name': 'test-schedule',
+    'model_name': 'test_model',
+    'model_id': 301,
+    'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 301',
+    'schema': 'predictions_schema',
+    'table_name': 'predictions',
+    'transform_table_name': 'transformed_data',
+    'input_filters': {
+        'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
+    },
+    'mlflow_transform_filters': {
+        'API_ERROR': {'policy': 'STOP', 'config': {}},
+    },
+    'mlflow_predict_filters': {
+        'API_ERROR': {'policy': 'STOP', 'config': {}},
+    },
+    'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    'opc_output_config': {},
+    'pi_web_api_output_config': {},
+    'save_transform': True,
+    'prediction_store_policy': 'lts:1',
+    'model_config': {
+        'retention_minutes': 0,
+        'transform_flavor': 'sklearn',
+        'predict_flavor': 'sklearn',
+    },
+    'datetime_columns': ['timestamp', 'created_at'],
+}
+
+base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}"
+
+def get_base_input_data(model_id):
+    return {
+        **base_input_data,
+        'model_id': model_id,
+        'query': base_query.format(model_id=model_id),
+    }
+
+def insert_sample_data(postgres_engine, model_id, values: list):
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}"))
+        
+        values_sql = []
+        for i, value in enumerate(values):
+            values_sql.append(f"""
+                ({model_id}, 'sensor_{i+1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
+            """)
+
+        insert_sql = f"""
+            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
+            VALUES
+                {', '.join(values_sql)}
+        """
+        conn.execute(text(insert_sql))
+
+async def start_and_await_workflow(client, input_data, workflow_id):
+    handle = await client.start_workflow(
+        PredictionsBatch.run,
+        input_data,
+        id=workflow_id,
+        task_queue='test-queue',
+    )
+    print("[TEST] ✓ Workflow started")
+
+    print("\n[TEST] 3. Waiting for workflow completion...")
+    try:
+        await asyncio.wait_for(handle.result(), timeout=60.0)
+        print("[TEST] ✓ Workflow completed successfully")
+    except asyncio.TimeoutError:
+        pytest.fail("Workflow execution timed out after 60 seconds")
+
 
 @pytest.mark.asyncio
 @pytest.mark.integration
@@ -46,122 +119,168 @@ async def test_scenario_3_1_1_default_prediction_export(
     """
     client = temporal_test_env.client
 
+    model_id = 311
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 301"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (301, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (301, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
-    # Mock mlflow_response_gate for predict to return a non-None path_flag
+    input_data = get_base_input_data(model_id)
-    # Any non-None path_flag that's not STOP/CONTINUE/REPEAT will be passed to format_and_export_prediction
+    input_data['pi_web_api_output_config'] = {
-    # which will then call format_default_prediction
+        'endpoint': 'test_endpoint',
-    from unittest.mock import patch
+        'prediction_tags': {'tag_1': 'web_id_1'},
-    
+        'confidence_tags': {'tag_2': 'web_id_2'},
-    original_mlflow_response_gate = test_activities.mlflow_response_gate
+    }
-    
+    input_data['opc_output_config'] = {
-    async def mock_mlflow_response_gate(input_data):
+        '1': {
-        # Only return error path_flag for predict, not transform
+            'prediction_tags': {
-        if input_data.get('type') == 'predict':
+                'addr_1': {
-            metadata = input_data.get('metadata', {})
+                    'data_type': 'float',
-            # Return a path_flag that will be passed to format_and_export_prediction
-            # but won't trigger early exit (not STOP/CONTINUE/REPEAT)
-            # The path_flag_handler only returns True for STOP/CONTINUE/REPEAT
-            # So any other value will make it return False and continue to export
-            return 'ERROR', -1, 'Error: Prediction validation failed'
-        # For transform, return normal (None)
-        return await original_mlflow_response_gate(input_data)
-    
-    with patch.object(test_activities, 'mlflow_response_gate', side_effect=mock_mlflow_response_gate):
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 301,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
                 }
             },
-            'schedule_name': 'test-schedule',
+            'confidence_tags': {
-            'model_name': 'test_model',
+                'addr_2': {
-            'model_id': 301,
+                    'data_type': 'float',
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 301',
+                }
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
             },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'CONTINUE', 'config': {}},  # CONTINUE allows workflow to proceed
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
         }
+    }
 
-        print("\n[TEST] 2. Starting workflow that should create default prediction...")
+    print("\n[TEST] 2. Starting workflow that should create default prediction...")
-        workflow_id = f'test-default-prediction-{datetime.now().timestamp()}'
+    workflow_id = f'test-default-prediction-{datetime.now().timestamp()}'
-        
+    
-        handle = await client.start_workflow(
+    await start_and_await_workflow(client, input_data, workflow_id)
-            PredictionsBatch.run,
-            input_data,
-            id=workflow_id,
-            task_queue='test-queue',
-        )
-        print("[TEST] ✓ Workflow started")
 
-        print("\n[TEST] 3. Waiting for workflow completion...")
-        try:
-            await asyncio.wait_for(handle.result(), timeout=60.0)
-            print("[TEST] ✓ Workflow completed successfully")
-        except asyncio.TimeoutError:
-            pytest.fail("Workflow execution timed out after 60 seconds")
 
-        print("\n[TEST] 4. Verifying default prediction was created...")
-        with postgres_engine.connect() as conn:
-            result_query = conn.execute(
-                text("SELECT prediction, prediction_confidence, comments FROM predictions_schema.predictions WHERE model_id = 301")
-            )
-            prediction_rows = result_query.fetchall()
-            # Should have a default prediction with error comment
-            assert len(prediction_rows) >= 1, "Expected at least one default prediction"
-            if len(prediction_rows) > 0:
-                row = prediction_rows[0]
-                # Default predictions typically have specific characteristics
-                # The exact values depend on format_default_prediction implementation
-                print(f"[TEST] Default prediction found: {row}")
-        
-        print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
 @pytest.mark.integration
-async def test_scenario_3_1_2_export_without_optional_outputs(
+async def test_scenario_3_1_2_export_with_opc_only(
     temporal_test_env: WorkflowEnvironment,
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
 ):
     """
-    Scenario 3.1.2: Export Without Optional Outputs
+    Scenario 3.1.2: Export with OPC only
+    
+    Description:
+    Export to PostgreSQL and OPC server only (no PI Web API).
+    
+    Expected Behavior:
+    - Normal formatting
+    - PostgreSQL export executed
+    - OPC export executed
+    - PI Web API activity skipped
+    - Metrics written with OPC metrics
+    
+    Assertions:
+    - PI Web API activity NOT called
+    - OPC activity called
+    - PostgreSQL export called
+    - Metrics written with opc_metrics populated
+    """
+    client = temporal_test_env.client
+
+    model_id = 312
+
+    print("\n[TEST] 1. Inserting test data...")
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+    print("[TEST] ✓ Data inserted successfully")
+
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = {
+        'server_name': 'test_server',
+        'tags': {'prediction': 'test_tag'},
+    }
+    input_data['pi_web_api_output_config'] = None  # No PI Web API config
+
+    print("\n[TEST] 2. Starting workflow with OPC only...")
+    workflow_id = f'test-opc-only-{datetime.now().timestamp()}'
+    
+    await start_and_await_workflow(client, input_data, workflow_id)
+
+    print("\n[TEST] 4. Verifying PostgreSQL and OPC export were executed...")
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
+        )
+        prediction_rows = result_query.fetchall()
+        assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
+    
+    print("\n[TEST] ✓ All assertions passed!")
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_3_export_with_pi_web_api_only(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.3: Export with PI Web API only
+    
+    Description:
+    Export to PostgreSQL and PI Web API only (no OPC).
+    
+    Expected Behavior:
+    - Normal formatting
+    - PostgreSQL export executed
+    - PI Web API export executed
+    - OPC activity skipped
+    - Metrics written without OPC metrics
+    
+    Assertions:
+    - OPC activity NOT called
+    - PI Web API activity called
+    - PostgreSQL export called
+    - Metrics written with empty opc_metrics
+    """
+    client = temporal_test_env.client
+
+    model_id = 313
+
+    print("\n[TEST] 1. Inserting test data...")
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+    print("[TEST] ✓ Data inserted successfully")
+
+    input_data = get_base_input_data(model_id)
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'prediction': 'test_pred_tag'},
+        'confidence_tags': {'confidence': 'test_conf_tag'},
+    }
+    input_data['opc_output_config'] = None  # No OPC config
+
+    print("\n[TEST] 2. Starting workflow with PI Web API only...")
+    workflow_id = f'test-pi-api-only-{datetime.now().timestamp()}'
+    
+    await start_and_await_workflow(client, input_data, workflow_id)
+
+    print("\n[TEST] 4. Verifying PostgreSQL and PI Web API export were executed...")
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
+        )
+        prediction_rows = result_query.fetchall()
+        assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
+    
+    print("\n[TEST] ✓ All assertions passed!")
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_4_export_without_optional_outputs(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.4: Export Without Optional Outputs
     
     Description:
     Export only to PostgreSQL (no OPC or PI Web API).
@@ -180,79 +299,25 @@ async def test_scenario_3_1_2_export_without_optional_outputs(
     """
     client = temporal_test_env.client
 
+    model_id = 314
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 302"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (302, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (302, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
-    input_data = {
+    input_data = get_base_input_data(model_id)
-        'metadata': {
+    input_data['opc_output_config'] = None  # No OPC config
-            'metadata': {
+    input_data['pi_web_api_output_config'] = None  # No PI Web API config
-                'model_id': 302,
-                'model_name': 'test_model',
-                'schedule_name': 'test-schedule',
-                'workflow_name': 'predictions_batch',
-            }
-        },
-        'schedule_name': 'test-schedule',
-        'model_name': 'test_model',
-        'model_id': 302,
-        'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 302',
-        'schema': 'predictions_schema',
-        'table_name': 'predictions',
-        'transform_table_name': 'transformed_data',
-        'input_filters': {
-            'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-        },
-        'mlflow_transform_filters': {
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
-        },
-        'mlflow_predict_filters': {
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
-        },
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-        'opc_output_config': None,  # No OPC config
-        'pi_web_api_output_config': None,  # No PI Web API config
-        'save_transform': True,
-        'prediction_store_policy': 'lts:1',
-        'model_config': {
-            'retention_minutes': 0,
-            'transform_flavor': 'sklearn',
-            'predict_flavor': 'sklearn',
-        },
-        'datetime_columns': ['timestamp', 'created_at'],
-    }
 
     print("\n[TEST] 2. Starting workflow without optional outputs...")
     workflow_id = f'test-no-optional-outputs-{datetime.now().timestamp()}'
     
-    handle = await client.start_workflow(
+    await start_and_await_workflow(client, input_data, workflow_id)
-        PredictionsBatch.run,
-        input_data,
-        id=workflow_id,
-        task_queue='test-queue',
-    )
-    print("[TEST] ✓ Workflow started")
-
-    print("\n[TEST] 3. Waiting for workflow completion...")
-    try:
-        await asyncio.wait_for(handle.result(), timeout=60.0)
-        print("[TEST] ✓ Workflow completed successfully")
-    except asyncio.TimeoutError:
-        pytest.fail("Workflow execution timed out after 60 seconds")
 
     print("\n[TEST] 4. Verifying only PostgreSQL export was executed...")
     with postgres_engine.connect() as conn:
         result_query = conn.execute(
-            text("SELECT model_id FROM predictions_schema.predictions WHERE model_id = 302")
+            text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
         )
         prediction_rows = result_query.fetchall()
         assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
@@ -262,14 +327,14 @@ async def test_scenario_3_1_2_export_without_optional_outputs(
 
 @pytest.mark.asyncio
 @pytest.mark.integration
-async def test_scenario_3_1_3_export_without_transformed_data(
+async def test_scenario_3_1_5_export_without_transformed_data(
     temporal_test_env: WorkflowEnvironment,
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
 ):
     """
-    Scenario 3.1.3: Export Without Transformed Data
+    Scenario 3.1.5: Export Without Transformed Data
     
     Description:
     Only prediction exported, no transform table.
@@ -286,88 +351,34 @@ async def test_scenario_3_1_3_export_without_transformed_data(
     """
     client = temporal_test_env.client
 
+    model_id = 315
+
     print("\n[TEST] 1. Inserting test data...")
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
     with postgres_engine.begin() as conn:
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 303"))
+        conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
-        conn.execute(text("DELETE FROM predictions_schema.transformed_data WHERE model_id = 303"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (303, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (303, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
-    input_data = {
+    input_data = get_base_input_data(model_id)
-        'metadata': {
+    input_data['save_transform'] = False  # Don't save transformed data
-            'metadata': {
-                'model_id': 303,
-                'model_name': 'test_model',
-                'schedule_name': 'test-schedule',
-                'workflow_name': 'predictions_batch',
-            }
-        },
-        'schedule_name': 'test-schedule',
-        'model_name': 'test_model',
-        'model_id': 303,
-        'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 303',
-        'schema': 'predictions_schema',
-        'table_name': 'predictions',
-        'transform_table_name': 'transformed_data',
-        'input_filters': {
-            'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-        },
-        'mlflow_transform_filters': {
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
-        },
-        'mlflow_predict_filters': {
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
-        },
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-        'opc_output_config': {},
-        'pi_web_api_output_config': {},
-        'save_transform': False,  # Don't save transformed data
-        'prediction_store_policy': 'lts:1',
-        'model_config': {
-            'retention_minutes': 0,
-            'transform_flavor': 'sklearn',
-            'predict_flavor': 'sklearn',
-        },
-        'datetime_columns': ['timestamp', 'created_at'],
-    }
 
     print("\n[TEST] 2. Starting workflow without transformed data export...")
     workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}'
     
-    handle = await client.start_workflow(
+    await start_and_await_workflow(client, input_data, workflow_id)
-        PredictionsBatch.run,
-        input_data,
-        id=workflow_id,
-        task_queue='test-queue',
-    )
-    print("[TEST] ✓ Workflow started")
-
-    print("\n[TEST] 3. Waiting for workflow completion...")
-    try:
-        await asyncio.wait_for(handle.result(), timeout=60.0)
-        print("[TEST] ✓ Workflow completed successfully")
-    except asyncio.TimeoutError:
-        pytest.fail("Workflow execution timed out after 60 seconds")
 
     print("\n[TEST] 4. Verifying only prediction was exported...")
     with postgres_engine.connect() as conn:
         # Verify prediction exists
         result_query = conn.execute(
-            text("SELECT model_id FROM predictions_schema.predictions WHERE model_id = 303")
+            text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
         )
         prediction_rows = result_query.fetchall()
         assert len(prediction_rows) == 1, "Expected one prediction record"
         
         # Verify transformed data table is empty
         result_query = conn.execute(
-            text("SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = 303")
+            text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
         )
         count = result_query.scalar()
         assert count == 0, f"Expected transform table to be empty, but found {count} records"
@@ -402,21 +413,13 @@ async def test_scenario_3_2_1_postgres_export_error_predictions_table(
     """
     client = temporal_test_env.client
 
+    model_id = 321
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 304"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (304, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (304, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
     # Mock export_data_to_postgres to raise an exception
-    from unittest.mock import patch
     original_export = test_activities.export_data_to_postgres
     call_count = {'count': 0}
     
@@ -428,43 +431,7 @@ async def test_scenario_3_2_1_postgres_export_error_predictions_table(
         return await original_export(*args, **kwargs)
     
     with patch.object(test_activities, 'export_data_to_postgres', side_effect=mock_export_data_to_postgres):
-        input_data = {
+        input_data = get_base_input_data(model_id)
-            'metadata': {
-                'metadata': {
-                    'model_id': 304,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 304,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 304',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
-        }
 
         print("\n[TEST] 2. Starting workflow that should fail on PostgreSQL export...")
         workflow_id = f'test-postgres-error-{datetime.now().timestamp()}'
@@ -486,7 +453,7 @@ async def test_scenario_3_2_1_postgres_export_error_predictions_table(
             # Verify no predictions were created
             with postgres_engine.connect() as conn:
                 result_query = conn.execute(
-                    text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 304")
+                    text(f"SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}")
                 )
                 count = result_query.scalar()
                 assert count == 0, f"Expected no predictions, but found {count} records"
@@ -521,65 +488,22 @@ async def test_scenario_3_2_2_pi_web_api_write_error(
     """
     client = temporal_test_env.client
 
+    model_id = 322
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 305"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (305, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (305, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
     # Mock write_pi_web_api_data to raise an exception
-    from unittest.mock import patch
     def mock_write_pi_web_api_data(*args, **kwargs):
         raise Exception("PI Web API service unavailable")
     
     with patch.object(test_activities, 'write_pi_web_api_data', side_effect=mock_write_pi_web_api_data):
-        input_data = {
+        input_data = get_base_input_data(model_id)
-            'metadata': {
+        input_data['pi_web_api_output_config'] = {
-                'metadata': {
+            'endpoint': 'test_endpoint',
-                    'model_id': 305,
+            'prediction_tags': {'prediction': 'test_pred_tag'},
-                    'model_name': 'test_model',
+            'confidence_tags': {'confidence': 'test_conf_tag'},
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 305,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 305',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {},
-            'pi_web_api_output_config': {
-                'endpoint': 'test_endpoint',
-                'prediction_tags': {'prediction': 'test_pred_tag'},
-                'confidence_tags': {'confidence': 'test_conf_tag'},
-            },
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
         }
 
         print("\n[TEST] 2. Starting workflow that should fail on PI Web API write...")
@@ -629,64 +553,21 @@ async def test_scenario_3_2_3_opc_write_error(
     """
     client = temporal_test_env.client
 
+    model_id = 323
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 306"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (306, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (306, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
     # Mock write_opc_data to raise an exception
-    from unittest.mock import patch
     def mock_write_opc_data(*args, **kwargs):
         raise Exception("OPC server unavailable")
     
     with patch.object(test_activities, 'write_opc_data', side_effect=mock_write_opc_data):
-        input_data = {
+        input_data = get_base_input_data(model_id)
-            'metadata': {
+        input_data['opc_output_config'] = {
-                'metadata': {
+            'server_name': 'test_server',
-                    'model_id': 306,
+            'tags': {'prediction': 'test_tag'},
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 306,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 306',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {
-                'server_name': 'test_server',
-                'tags': {'prediction': 'test_tag'},
-            },
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
         }
 
         print("\n[TEST] 2. Starting workflow that should fail on OPC write...")

e2e/test_predictions_batch_prediction_process.py

@@ -407,6 +407,8 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(
 
     print("\n[TEST] ✓ All assertions passed!")
 
+
+
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_2_2_3_transform_gate_triggers_repeat(
@@ -414,6 +416,7 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat(
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
+    bad_data_model,
 ):
     """
     Scenario 2.2.3: Transform Gate Triggers REPEAT
@@ -456,6 +459,28 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat(
 
     print("\n[TEST] ✓ All assertions passed!")
 
+
+@pytest.fixture
+def bad_predict_model(
+    patch_mlflow,
+    mock_mlflow_models
+):
+    model = MagicMock(
+        predict=MagicMock(
+            side_effect=Exception("Bad predict model")
+        )
+    )
+
+    def mock_sklearn_load_model(model_uri):
+        if 'data_model' in model_uri or 'transform' in model_uri.lower():
+            return mock_mlflow_models['transform_model']
+        return model
+    patch_mlflow.sklearn = MagicMock()
+    patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
+
+    return model
+
+
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_2_3_1_predict_gate_triggers_continue(
@@ -463,6 +488,7 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue(
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
+    bad_predict_model,
 ):
     """
     Scenario 2.3.1: Predict Gate Triggers CONTINUE
@@ -484,93 +510,26 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue(
     """
     client = temporal_test_env.client
 
+    model_id = 231
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 210"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (210, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (210, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
-    print("[TEST] ✓ Data inserted successfully")
 
-    # Mock mlflow_response_gate for predict to return CONTINUE
+    input_data = get_base_input_data(model_id)
-    original_mlflow_response_gate = test_activities.mlflow_response_gate
+    input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'CONTINUE'
+
+    print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...")
+    workflow_id = f'test-predict-continue-{datetime.now().timestamp()}'
     
-    async def mock_mlflow_response_gate(input_data):
+    await start_and_await_workflow(client, input_data, workflow_id)
-        if input_data.get('type') == 'predict':
+    assert_continue(
-            return 'CONTINUE', 10, 'Predict response has issues but continuing'
+        postgres_engine=postgres_engine,
-        return await original_mlflow_response_gate(input_data)
+        model_id=model_id,
+        prediction_confidence=Decimal(10),
+        comments='Bad predict model',
+    )
     
-    with patch.object(test_activities, 'mlflow_response_gate', side_effect=mock_mlflow_response_gate):
+    print("\n[TEST] ✓ All assertions passed!")
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 210,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 210,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 210',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'CONTINUE', 'config': {}},  # CONTINUE on predict error
-            },
-            'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],  # CONTINUE first
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
-        }
-
-        print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...")
-        workflow_id = f'test-predict-continue-{datetime.now().timestamp()}'
-        
-        handle = await client.start_workflow(
-            PredictionsBatch.run,
-            input_data,
-            id=workflow_id,
-            task_queue='test-queue',
-        )
-        print("[TEST] ✓ Workflow started")
-
-        print("\n[TEST] 3. Waiting for workflow completion...")
-        try:
-            await asyncio.wait_for(handle.result(), timeout=60.0)
-            print("[TEST] ✓ Workflow completed successfully")
-        except asyncio.TimeoutError:
-            pytest.fail("Workflow execution timed out after 60 seconds")
-
-        print("\n[TEST] 4. Verifying prediction was created (via CONTINUE path)...")
-        with postgres_engine.connect() as conn:
-            result_query = conn.execute(
-                text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 210")
-            )
-            count = result_query.scalar()
-            assert count >= 1, f"Expected at least one prediction, but found {count} records"
-        
-        print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
@@ -580,6 +539,7 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop(
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
+    bad_predict_model,
 ):
     """
     Scenario 2.3.2: Predict Gate Triggers STOP
@@ -600,93 +560,22 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop(
     """
     client = temporal_test_env.client
 
+    model_id = 232
+
     print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 205"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (205, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (205, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
     print("[TEST] ✓ Data inserted successfully")
 
-    # Mock mlflow_response_gate for predict to return STOP
+    input_data = get_base_input_data(model_id)
-    original_mlflow_response_gate = test_activities.mlflow_response_gate
+    input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'STOP'
+
+    print("\n[TEST] 2. Starting workflow that should stop at predict gate...")
+    workflow_id = f'test-predict-stop-{datetime.now().timestamp()}'
     
-    async def mock_mlflow_response_gate(input_data):
+    await start_and_await_workflow(client, input_data, workflow_id)
-        if input_data.get('type') == 'predict':
+    assert_stop(postgres_engine, model_id)
-            return 'STOP', -1, 'Predict response validation failed'
-        return await original_mlflow_response_gate(input_data)
-    
-    with patch.object(test_activities, 'mlflow_response_gate', side_effect=mock_mlflow_response_gate):
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 205,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 205,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 205',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},  # STOP on predict error
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
-        }
 
-        print("\n[TEST] 2. Starting workflow that should stop at predict gate...")
+    print("\n[TEST] ✓ All assertions passed!")
-        workflow_id = f'test-predict-stop-{datetime.now().timestamp()}'
-        
-        handle = await client.start_workflow(
-            PredictionsBatch.run,
-            input_data,
-            id=workflow_id,
-            task_queue='test-queue',
-        )
-        print("[TEST] ✓ Workflow started")
-
-        print("\n[TEST] 3. Waiting for workflow completion (should exit early)...")
-        try:
-            await asyncio.wait_for(handle.result(), timeout=60.0)
-            print("[TEST] ✓ Workflow completed (exited early as expected)")
-        except asyncio.TimeoutError:
-            pytest.fail("Workflow execution timed out after 60 seconds")
-
-        print("\n[TEST] 4. Verifying no predictions were created...")
-        with postgres_engine.connect() as conn:
-            result_query = conn.execute(
-                text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 205")
-            )
-            count = result_query.scalar()
-            assert count == 0, f"Expected no predictions, but found {count} records"
-        
-        print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
@@ -696,6 +585,7 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat(
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
+    bad_predict_model,
 ):
     """
     Scenario 2.3.3: Predict Gate Triggers REPEAT
@@ -718,323 +608,22 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat(
     """
     client = temporal_test_env.client
 
+    model_id = 233
+
     print("\n[TEST] 1. Inserting test data and previous prediction...")
-    with postgres_engine.begin() as conn:
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 211"))
+    data = insert_sample_prediction(postgres_engine, model_id)
-        conn.execute(text("DELETE FROM predictions_schema.predictions WHERE model_id = 211"))
+
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (211, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (211, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
-        
-        # Insert a previous prediction to repeat
-        insert_prediction = """
-            INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
-            VALUES (211, '2024-01-01 11:00:00+00:00', 47.5, 10, 'Good', 'Previous prediction', 0.1)
-        """
-        conn.execute(text(insert_prediction))
     print("[TEST] ✓ Data and previous prediction inserted")
 
-    # Mock request_predict to return an error response
+    input_data = get_base_input_data(model_id)
-    def mock_request_predict(*args, **kwargs):
+    input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'REPEAT'
-        return {
+    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']  # REPEAT first
-            'success': False,  # This will trigger API_ERROR filter
-            'content': [],
-        }
-    
-    with patch.object(test_activities, 'request_predict', side_effect=mock_request_predict):
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 211,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 211,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 211',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'REPEAT', 'config': {}},  # REPEAT on predict error
-            },
-            'path_priority': ['REPEAT', 'STOP', 'CONTINUE'],  # REPEAT first
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
-        }
 
-        print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...")
+    print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...")
-        workflow_id = f'test-predict-repeat-{datetime.now().timestamp()}'
+    workflow_id = f'test-predict-repeat-{datetime.now().timestamp()}'
-        
-        handle = await client.start_workflow(
-            PredictionsBatch.run,
-            input_data,
-            id=workflow_id,
-            task_queue='test-queue',
-        )
-        print("[TEST] ✓ Workflow started")
 
-        print("\n[TEST] 3. Waiting for workflow completion...")
+    await start_and_await_workflow(client, input_data, workflow_id)
-        try:
+    assert_repeat(postgres_engine, model_id, data)
-            await asyncio.wait_for(handle.result(), timeout=60.0)
-            print("[TEST] ✓ Workflow completed successfully")
-        except asyncio.TimeoutError:
-            pytest.fail("Workflow execution timed out after 60 seconds")
 
-        print("\n[TEST] 4. Verifying prediction was repeated...")
+    print("\n[TEST] ✓ All assertions passed!")
-        with postgres_engine.connect() as conn:
-            result_query = conn.execute(
-                text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 211")
-            )
-            count = result_query.scalar()
-            assert count >= 2, f"Expected at least 2 predictions (original + repeated), but found {count} records"
-        
-        print("\n[TEST] ✓ All assertions passed!")
-
-
-@pytest.mark.asyncio
-@pytest.mark.integration
-async def test_scenario_2_4_1_mlflow_transform_api_error(
-    temporal_test_env: WorkflowEnvironment,
-    temporal_worker: Worker,
-    test_activities: Activities,
-    postgres_engine,
-):
-    """
-    Scenario 2.4.1: MLFlow Transform API Error
-    
-    Description:
-    MLFlow transform request fails.
-    
-    Expected Behavior:
-    - request_transform raises exception
-    - Notification sent with MLFlow error details
-    - Workflow fails after retry attempts exhausted
-    
-    Assertions:
-    - Exception raised from transform activity
-    - Error notification sent
-    - Workflow fails
-    - Export NOT called
-    """
-    client = temporal_test_env.client
-
-    print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 207"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (207, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (207, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
-    print("[TEST] ✓ Data inserted successfully")
-
-    # Mock request_transform to raise an exception
-    def mock_request_transform(*args, **kwargs):
-        raise Exception("MLFlow transform service unavailable")
-    
-    with patch.object(test_activities, 'request_transform', side_effect=mock_request_transform):
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 207,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 207,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 207',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
-        }
-
-        print("\n[TEST] 2. Starting workflow that should fail on transform...")
-        workflow_id = f'test-transform-error-{datetime.now().timestamp()}'
-        
-        handle = await client.start_workflow(
-            PredictionsBatch.run,
-            input_data,
-            id=workflow_id,
-            task_queue='test-queue',
-        )
-        print("[TEST] ✓ Workflow started")
-
-        print("\n[TEST] 3. Waiting for workflow completion...")
-        try:
-            await asyncio.wait_for(handle.result(), timeout=60.0)
-            print("[TEST] ✓ Workflow completed (may have failed after retries or completed with error handling)")
-        except Exception as e:
-            print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
-        
-        # Verify no predictions were created (regardless of whether workflow failed or completed)
-        print("\n[TEST] 4. Verifying no predictions were created...")
-        with postgres_engine.connect() as conn:
-            result_query = conn.execute(
-                text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 207")
-            )
-            count = result_query.scalar()
-            assert count == 0, f"Expected no predictions, but found {count} records"
-        
-        print("\n[TEST] ✓ All assertions passed!")
-
-
-@pytest.mark.asyncio
-@pytest.mark.integration
-async def test_scenario_2_4_2_mlflow_predict_api_error(
-    temporal_test_env: WorkflowEnvironment,
-    temporal_worker: Worker,
-    test_activities: Activities,
-    postgres_engine,
-):
-    """
-    Scenario 2.4.2: MLFlow Predict API Error
-    
-    Description:
-    MLFlow predict request fails.
-    
-    Expected Behavior:
-    - request_predict raises exception
-    - Notification sent
-    - Workflow fails after retries
-    
-    Assertions:
-    - Transform succeeded
-    - Predict raised exception
-    - Error notification sent
-    - Workflow fails
-    """
-    client = temporal_test_env.client
-
-    print("\n[TEST] 1. Inserting test data...")
-    with postgres_engine.begin() as conn:
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 208"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (208, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (208, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-        """
-        conn.execute(text(insert_sql))
-    print("[TEST] ✓ Data inserted successfully")
-
-    # Mock request_predict to raise an exception
-    def mock_request_predict(*args, **kwargs):
-        raise Exception("MLFlow predict service unavailable")
-    
-    with patch.object(test_activities, 'request_predict', side_effect=mock_request_predict):
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 208,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 208,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 208',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-            'opc_output_config': {},
-            'pi_web_api_output_config': {},
-            'save_transform': True,
-            'prediction_store_policy': 'lts:1',
-            'model_config': {
-                'retention_minutes': 0,
-                'transform_flavor': 'sklearn',
-                'predict_flavor': 'sklearn',
-            },
-            'datetime_columns': ['timestamp', 'created_at'],
-        }
-
-        print("\n[TEST] 2. Starting workflow that should fail on predict...")
-        workflow_id = f'test-predict-error-{datetime.now().timestamp()}'
-        
-        handle = await client.start_workflow(
-            PredictionsBatch.run,
-            input_data,
-            id=workflow_id,
-            task_queue='test-queue',
-        )
-        print("[TEST] ✓ Workflow started")
-
-        print("\n[TEST] 3. Waiting for workflow completion...")
-        try:
-            await asyncio.wait_for(handle.result(), timeout=60.0)
-            print("[TEST] ✓ Workflow completed (may have failed after retries or completed with error handling)")
-        except Exception as e:
-            print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
-        
-        # Verify no predictions were created (regardless of whether workflow failed or completed)
-        print("\n[TEST] 4. Verifying no predictions were created...")
-        with postgres_engine.connect() as conn:
-            result_query = conn.execute(
-                text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 208")
-            )
-            count = result_query.scalar()
-            assert count == 0, f"Expected no predictions, but found {count} records"
-        
-        print("\n[TEST] ✓ All assertions passed!")