diff --git a/e2e/scenarios.md b/e2e/scenarios.md index f9f8bbc..efdb7bb 100644 --- a/e2e/scenarios.md +++ b/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 diff --git a/e2e/test_predictions_batch_format_export.py b/e2e/test_predictions_batch_format_export.py index 901d15a..c5512c2 100644 --- a/e2e/test_predictions_batch_format_export.py +++ b/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: - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) print("[TEST] ✓ Data inserted successfully") - # Mock mlflow_response_gate for predict to return a non-None path_flag - # Any non-None path_flag that's not STOP/CONTINUE/REPEAT will be passed to format_and_export_prediction - # which will then call format_default_prediction - from unittest.mock import patch - - original_mlflow_response_gate = test_activities.mlflow_response_gate - - async def mock_mlflow_response_gate(input_data): - # Only return error path_flag for predict, not transform - if input_data.get('type') == 'predict': - metadata = input_data.get('metadata', {}) - # 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', + input_data = get_base_input_data(model_id) + input_data['pi_web_api_output_config'] = { + 'endpoint': 'test_endpoint', + 'prediction_tags': {'tag_1': 'web_id_1'}, + 'confidence_tags': {'tag_2': 'web_id_2'}, + } + input_data['opc_output_config'] = { + '1': { + 'prediction_tags': { + 'addr_1': { + 'data_type': 'float', } }, - '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': {}}, + 'confidence_tags': { + 'addr_2': { + 'data_type': 'float', + } }, - '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...") - workflow_id = f'test-default-prediction-{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] 2. Starting workflow that should create default prediction...") + workflow_id = f'test-default-prediction-{datetime.now().timestamp()}' + + await start_and_await_workflow(client, input_data, workflow_id) - 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: - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) print("[TEST] ✓ Data inserted successfully") - input_data = { - 'metadata': { - 'metadata': { - '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'], - } + input_data = get_base_input_data(model_id) + input_data['opc_output_config'] = None # No OPC config + input_data['pi_web_api_output_config'] = None # No PI Web API config print("\n[TEST] 2. Starting workflow without optional outputs...") workflow_id = f'test-no-optional-outputs-{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") + await start_and_await_workflow(client, input_data, workflow_id) 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("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)) + conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) print("[TEST] ✓ Data inserted successfully") - input_data = { - 'metadata': { - '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'], - } + input_data = get_base_input_data(model_id) + input_data['save_transform'] = False # Don't save transformed data 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( - 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") + await start_and_await_workflow(client, input_data, workflow_id) 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: - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) 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 = { - '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'], - } + input_data = get_base_input_data(model_id) 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: - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) 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 = { - 'metadata': { - 'metadata': { - 'model_id': 305, - 'model_name': 'test_model', - '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'], + 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'}, } 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: - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) 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 = { - 'metadata': { - 'metadata': { - 'model_id': 306, - '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'], + input_data = get_base_input_data(model_id) + input_data['opc_output_config'] = { + 'server_name': 'test_server', + 'tags': {'prediction': 'test_tag'}, } print("\n[TEST] 2. Starting workflow that should fail on OPC write...") diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index 0bd75da..d32980e 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/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: - 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") + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - # Mock mlflow_response_gate for predict to return CONTINUE - original_mlflow_response_gate = test_activities.mlflow_response_gate + input_data = get_base_input_data(model_id) + 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): - if input_data.get('type') == 'predict': - return 'CONTINUE', 10, 'Predict response has issues but continuing' - return await original_mlflow_response_gate(input_data) + 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 predict model', + ) - with patch.object(test_activities, 'mlflow_response_gate', side_effect=mock_mlflow_response_gate): - 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!") + 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: - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) print("[TEST] ✓ Data inserted successfully") - # Mock mlflow_response_gate for predict to return STOP - original_mlflow_response_gate = test_activities.mlflow_response_gate + input_data = get_base_input_data(model_id) + 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): - if input_data.get('type') == 'predict': - 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'], - } + await start_and_await_workflow(client, input_data, workflow_id) + assert_stop(postgres_engine, model_id) - print("\n[TEST] 2. Starting workflow that should stop at predict gate...") - 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!") + 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: - conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 211")) - 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)) + insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) + data = insert_sample_prediction(postgres_engine, model_id) + print("[TEST] ✓ Data and previous prediction inserted") - # Mock request_predict to return an error response - def mock_request_predict(*args, **kwargs): - return { - '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'], - } + input_data = get_base_input_data(model_id) + input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'REPEAT' + input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE'] # REPEAT first - print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...") - 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] 2. Starting workflow that should trigger REPEAT at predict gate...") + workflow_id = f'test-predict-repeat-{datetime.now().timestamp()}' - 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") + await start_and_await_workflow(client, input_data, workflow_id) + assert_repeat(postgres_engine, model_id, data) - 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 = 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!") + print("\n[TEST] ✓ All assertions passed!")