gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1478

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2026-01-16 10:38:08 -03:00
parent 0cd6ae660a
commit 6b9bb38968
3 changed files with 374 additions and 898 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
e2e/scenarios.md
View File

@@ -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

627
e2e/test_predictions_batch_format_export.py
View File

@@ -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...")

547
e2e/test_predictions_batch_prediction_process.py
View File

@@ -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!")