From 0cd6ae660a83eb6885f07538502ac16bd8e28d3c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 15 Jan 2026 16:17:47 -0300 Subject: [PATCH] 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. --- ...st_predictions_batch_prediction_process.py | 289 +++--------------- 1 file changed, 43 insertions(+), 246 deletions(-) diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index 7179a2f..0bd75da 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/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: - 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") - - 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'], - } + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + + input_data = get_base_input_data(model_id) + input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'CONTINUE' 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( - PredictionsBatch.run, - input_data, - id=workflow_id, - task_queue='test-queue', + await start_and_await_workflow(client, input_data, workflow_id) + assert_continue( + postgres_engine=postgres_engine, + model_id=model_id, + 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: - 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)) + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) print("[TEST] ✓ Data inserted successfully") - # Mock mlflow_response_gate for transform to return STOP - original_mlflow_response_gate = test_activities.mlflow_response_gate + input_data = get_base_input_data(model_id) + 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): - if input_data.get('type') == 'transform': - 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!") + await start_and_await_workflow(client, input_data, workflow_id) + assert_stop(postgres_engine, model_id) + 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...") - 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") + model_id = 223 - # Mock request_transform to return an error response - def mock_request_transform(*args, **kwargs): - return { - '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] 1. Inserting test data...") + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + data = insert_sample_prediction(postgres_engine, model_id) - print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...") - 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("[TEST] ✓ Data inserted successfully") - 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") + input_data = get_base_input_data(model_id) + input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'REPEAT' - print("\n[TEST] 4. Verifying prediction was repeated...") - with postgres_engine.connect() as conn: - 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!") + print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...") + workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}' + 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