SIENTIAPDE-1478

Refactor end-to-end tests for PredictionsBatch workflow

- Updated the `assert_continue` function to accept dynamic prediction confidence and comments.
- Simplified test scenarios by introducing helper functions for data insertion.
- Enhanced test cases to verify behavior for CONTINUE, STOP, and REPEAT policies at transform gates.
- Improved clarity and maintainability of test structure.

e2e/test_predictions_batch_prediction_process.py

@@ -110,7 +110,10 @@ async def start_and_await_workflow(client, input_data, workflow_id):
     except asyncio.TimeoutError:
         pytest.fail("Workflow execution timed out after 60 seconds")
 
-def assert_continue(postgres_engine, model_id):
+def assert_continue(
+    postgres_engine, model_id, prediction_confidence: Decimal = 2,
+    comments: str = 'Input data with bad quality',
+):
     print("\n[TEST] 4. Verifying prediction was created despite warnings...")
     with postgres_engine.connect() as conn:
         result_query = conn.execute(
@@ -122,9 +125,9 @@ def assert_continue(postgres_engine, model_id):
         # Assert prediction value is 0 and other fields
         row = prediction_rows[0]
         assert row[1] == 0, f"Expected prediction=0, got {row[1]}"
-        assert row[2] == 2, f"Expected prediction_confidence=0, got {row[2]}"
+        assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}"
         assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
-        assert row[4] == 'Input data with bad quality', f"Expected comments='Input data with bad quality', got {row[4]}"
+        assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
 
 
 def assert_stop(postgres_engine, model_id):
@@ -336,82 +339,24 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue(
     """
     client = temporal_test_env.client
 
+    model_id = 221
+
     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 = 207"))
+
-        
+    input_data = get_base_input_data(model_id)
-        insert_sql = """
+    input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'CONTINUE'
-            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")
-    
-    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': 'CONTINUE', 'config': {}},  # CONTINUE on transform error
-        },
-        'mlflow_predict_filters': {
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
-        },
-        '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 transform gate...")
     workflow_id = f'test-transform-continue-{datetime.now().timestamp()}'
     
-    handle = await client.start_workflow(
+    await start_and_await_workflow(client, input_data, workflow_id)
-        PredictionsBatch.run,
+    assert_continue(
-        input_data,
+        postgres_engine=postgres_engine,
-        id=workflow_id,
+        model_id=model_id,
-        task_queue='test-queue',
+        prediction_confidence=Decimal(10),
+        comments='Bad data model',
     )
-    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 = 207")
-        )
-        count = result_query.scalar()
-        assert count >= 1, f"Expected at least one prediction, but found {count} records"
     
     print("\n[TEST] ✓ All assertions passed!")
 
@@ -423,6 +368,7 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(
     temporal_worker: Worker,
     test_activities: Activities,
     postgres_engine,
+    bad_data_model,
 ):
     """
     Scenario 2.2.2: Transform Gate Triggers STOP
@@ -444,94 +390,22 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(
     """
     client = temporal_test_env.client
 
+    model_id = 222
+
     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 = 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 mlflow_response_gate for transform to return STOP
+    input_data = get_base_input_data(model_id)
-    original_mlflow_response_gate = test_activities.mlflow_response_gate
+    input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'STOP'
+
+    print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...")
+    workflow_id = f'test-transform-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') == 'transform':
+    assert_stop(postgres_engine, model_id)
-            return 'STOP', -1, 'Transform 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': 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': {}},  # STOP on transform error
-            },
-            '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 stop at transform gate...")
-        workflow_id = f'test-transform-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 = 208")
-            )
-            count = result_query.scalar()
-            assert count == 0, f"Expected no predictions, but found {count} records"
-        
-        print("\n[TEST] ✓ All assertions passed!")
 
+    print("\n[TEST] ✓ All assertions passed!")
 
 @pytest.mark.asyncio
 @pytest.mark.integration
@@ -563,101 +437,24 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat(
     """
     client = temporal_test_env.client
 
-    print("\n[TEST] 1. Inserting test data and previous prediction...")
+    model_id = 223
-    with postgres_engine.begin() as conn:
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 209"))
-        conn.execute(text("DELETE FROM predictions_schema.predictions WHERE model_id = 209"))
-        
-        insert_sql = """
-            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
-                (209, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
-                (209, '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 (209, '2024-01-01 11:00:00+00:00', 45.0, 10, 'Good', 'Previous prediction', 0.1)
-        """
-        conn.execute(text(insert_prediction))
-    print("[TEST] ✓ Data and previous prediction inserted")
 
-    # Mock request_transform to return an error response
+    print("\n[TEST] 1. Inserting test data...")
-    def mock_request_transform(*args, **kwargs):
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
-        return {
+    data = insert_sample_prediction(postgres_engine, model_id)
-            'success': False,  # This will trigger API_ERROR filter
-            'content': [],
-        }
-    
-    with patch.object(test_activities, 'request_transform', side_effect=mock_request_transform):
-        input_data = {
-            'metadata': {
-                'metadata': {
-                    'model_id': 209,
-                    'model_name': 'test_model',
-                    'schedule_name': 'test-schedule',
-                    'workflow_name': 'predictions_batch',
-                }
-            },
-            'schedule_name': 'test-schedule',
-            'model_name': 'test_model',
-            'model_id': 209,
-            'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 209',
-            'schema': 'predictions_schema',
-            'table_name': 'predictions',
-            'transform_table_name': 'transformed_data',
-            'input_filters': {
-                'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
-            },
-            'mlflow_transform_filters': {
-                'API_ERROR': {'policy': 'REPEAT', 'config': {}},  # REPEAT on transform error
-            },
-            'mlflow_predict_filters': {
-                'API_ERROR': {'policy': 'STOP', 'config': {}},
-            },
-            '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 transform gate...")
+    print("[TEST] ✓ Data inserted successfully")
-        workflow_id = f'test-transform-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...")
+    input_data = get_base_input_data(model_id)
-        try:
+    input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'REPEAT'
-            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] 2. Starting workflow that should trigger REPEAT at transform gate...")
-        with postgres_engine.connect() as conn:
+    workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}'
-            result_query = conn.execute(
-                text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 209")
-            )
-            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!")
 
+    await start_and_await_workflow(client, input_data, workflow_id)
+    assert_repeat(postgres_engine, model_id, data)
+
+    print("\n[TEST] ✓ All assertions passed!")
 
 @pytest.mark.asyncio
 @pytest.mark.integration