vitor-aignosi
0cd6ae660a
Refactor end-to-end tests for PredictionsBatch workflow - Updated the `assert_continue` function to accept dynamic prediction confidence and comments. - Simplified test scenarios by introducing helper functions for data insertion. - Enhanced test cases to verify behavior for CONTINUE, STOP, and REPEAT policies at transform gates. - Improved clarity and maintainability of test structure.
1041 lines
38 KiB
Python
1041 lines
38 KiB
Python
"""
|
|
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import datetime
|
|
from decimal import Decimal
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from pytz import timezone
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
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': 201,
|
|
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'input_filters': {
|
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
|
'policy': 'CONTINUE', # Continue despite issues, not STOP
|
|
'config': {'variables': ['sensor_1']},
|
|
},
|
|
},
|
|
'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[tuple]):
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}"))
|
|
|
|
# Insert data with some null values (quality issue)
|
|
|
|
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))
|
|
|
|
|
|
def insert_sample_prediction(postgres_engine, model_id):
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}"))
|
|
|
|
# Insert data with some null values (quality issue)
|
|
|
|
insert_sql = f"""
|
|
INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
|
|
VALUES
|
|
({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1)
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
|
|
return (model_id, Decimal(10), Decimal(0), 'Good')
|
|
|
|
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")
|
|
|
|
def assert_continue(
|
|
postgres_engine, model_id, prediction_confidence: Decimal = 2,
|
|
comments: str = 'Input data with bad quality',
|
|
):
|
|
print("\n[TEST] 4. Verifying prediction was created despite warnings...")
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}")
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1, "Expected one prediction record despite warnings"
|
|
|
|
# Assert prediction value is 0 and other fields
|
|
row = prediction_rows[0]
|
|
assert row[1] == 0, f"Expected prediction=0, got {row[1]}"
|
|
assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}"
|
|
assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
|
|
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
|
|
|
|
|
def assert_stop(postgres_engine, model_id):
|
|
print("\n[TEST] 4. Verifying no predictions were created...")
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
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"
|
|
|
|
def assert_repeat(postgres_engine, model_id, last_prediction: list):
|
|
print("\n[TEST] 4. Verifying prediction was repeated...")
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(f'SELECT model_id, prediction, prediction_confidence, prediction_status FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
print(prediction_rows)
|
|
assert len(prediction_rows) == 2, "Expected two prediction records"
|
|
assert prediction_rows[0] == last_prediction, f"Expected first prediction to be the same as the last prediction, got {prediction_rows[0]}, expected {last_prediction}"
|
|
assert prediction_rows[1] == last_prediction, f"Expected second prediction to be the same as the last prediction, got {prediction_rows[1]}, expected {last_prediction}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_1_input_gate_triggers_continue(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.1.1: Input Gate Triggers CONTINUE
|
|
|
|
Description:
|
|
Input gate determines data should use previous prediction.
|
|
|
|
Expected Behavior:
|
|
- input_gate returns path_flag='CONTINUE'
|
|
- path_flag_handler calls export workflow with input data directly
|
|
- MLFlow transform and predict skipped
|
|
- Data exported as-is
|
|
|
|
Assertions:
|
|
- input_gate called
|
|
- MLFlow operations NOT called
|
|
- Export workflow called with original data
|
|
- Workflow completes
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
model_id = 211
|
|
|
|
print("\n[TEST] 1. Inserting test data...")
|
|
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
|
|
print("[TEST] ✓ Data inserted successfully")
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
|
|
|
|
print("\n[TEST] 2. Starting workflow with CONTINUE policy...")
|
|
workflow_id = f'test-continue-policy-{datetime.now().timestamp()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
assert_continue(postgres_engine, model_id)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_2_input_gate_triggers_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.1.2: Input Gate Triggers STOP
|
|
|
|
Description:
|
|
Input data quality gate fails with STOP policy.
|
|
|
|
Expected Behavior:
|
|
- input_gate returns path_flag='STOP'
|
|
- path_flag_handler detects STOP
|
|
- Workflow returns early without calling MLFlow
|
|
- No prediction exported
|
|
|
|
Assertions:
|
|
- input_gate called
|
|
- path_flag_handler returns True (early exit)
|
|
- MLFlow transform NOT called
|
|
- Export workflow NOT called
|
|
- Workflow completes without error
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
model_id = 212
|
|
|
|
print("\n[TEST] 1. Inserting test data...")
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
print("[TEST] ✓ Data inserted successfully")
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['policy'] = 'STOP'
|
|
|
|
print("\n[TEST] 2. Starting workflow that should stop at input gate...")
|
|
workflow_id = f'test-input-stop-{datetime.now().timestamp()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_3_input_gate_triggers_repeat(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.1.3: Input Gate Triggers REPEAT
|
|
|
|
Description:
|
|
Input gate determines data should repeat last prediction.
|
|
|
|
Expected Behavior:
|
|
- input_gate returns path_flag='REPEAT'
|
|
- path_flag_handler calls repeat_last_prediction activity
|
|
- MLFlow transform and predict skipped
|
|
- Last prediction repeated and exported
|
|
|
|
Assertions:
|
|
- input_gate called
|
|
- MLFlow operations NOT called
|
|
- repeat_last_prediction activity called
|
|
- Workflow completes
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
model_id = 213
|
|
|
|
print("\n[TEST] 1. Inserting test data and previous prediction...")
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
data = insert_sample_prediction(postgres_engine, model_id)
|
|
|
|
print("[TEST] ✓ Data and previous prediction inserted")
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['policy'] = 'REPEAT'
|
|
|
|
print("\n[TEST] 2. Starting workflow that should trigger REPEAT...")
|
|
workflow_id = f'test-input-repeat-{datetime.now().timestamp()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|
|
|
|
@pytest.fixture
|
|
def bad_data_model(patch_mlflow):
|
|
model = MagicMock(
|
|
predict=MagicMock(
|
|
side_effect=Exception("Bad data model")
|
|
)
|
|
)
|
|
|
|
patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
|
|
|
|
return model
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_1_transform_gate_triggers_continue(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
):
|
|
"""
|
|
Scenario 2.2.1: Transform Gate Triggers CONTINUE
|
|
|
|
Description:
|
|
Transform response gate determines data should continue despite issues.
|
|
|
|
Expected Behavior:
|
|
- request_transform succeeds
|
|
- mlflow_response_gate for transform returns path_flag='CONTINUE'
|
|
- path_flag_handler calls export workflow with transform data
|
|
- MLFlow predict skipped
|
|
- Transform data exported as-is
|
|
|
|
Assertions:
|
|
- Transform completed
|
|
- mlflow_response_gate called for transform
|
|
- MLFlow predict NOT called
|
|
- Export workflow called with transform data
|
|
- Workflow completes
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
model_id = 221
|
|
|
|
print("\n[TEST] 1. Inserting test data...")
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'CONTINUE'
|
|
|
|
print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at transform gate...")
|
|
workflow_id = f'test-transform-continue-{datetime.now().timestamp()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
assert_continue(
|
|
postgres_engine=postgres_engine,
|
|
model_id=model_id,
|
|
prediction_confidence=Decimal(10),
|
|
comments='Bad data model',
|
|
)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_2_transform_gate_triggers_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
):
|
|
"""
|
|
Scenario 2.2.2: Transform Gate Triggers STOP
|
|
|
|
Description:
|
|
Transform response validation fails with STOP policy.
|
|
|
|
Expected Behavior:
|
|
- request_transform succeeds but response invalid
|
|
- mlflow_response_gate for transform returns path_flag='STOP'
|
|
- Workflow exits without calling predict or export
|
|
|
|
Assertions:
|
|
- Transform completed but validation failed
|
|
- mlflow_response_gate called for transform
|
|
- MLFlow predict NOT called
|
|
- Export workflow NOT called
|
|
- Workflow completes without error
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
model_id = 222
|
|
|
|
print("\n[TEST] 1. Inserting test data...")
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
print("[TEST] ✓ Data inserted successfully")
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'STOP'
|
|
|
|
print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...")
|
|
workflow_id = f'test-transform-stop-{datetime.now().timestamp()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_3_transform_gate_triggers_repeat(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.2.3: Transform Gate Triggers REPEAT
|
|
|
|
Description:
|
|
Transform response gate determines data should repeat last prediction.
|
|
|
|
Expected Behavior:
|
|
- request_transform succeeds but response has issues
|
|
- mlflow_response_gate for transform returns path_flag='REPEAT'
|
|
- path_flag_handler calls repeat_last_prediction activity
|
|
- MLFlow predict skipped
|
|
- Last prediction repeated and exported
|
|
|
|
Assertions:
|
|
- Transform completed but validation triggered REPEAT
|
|
- mlflow_response_gate called for transform
|
|
- MLFlow predict NOT called
|
|
- repeat_last_prediction activity called
|
|
- Workflow completes
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
model_id = 223
|
|
|
|
print("\n[TEST] 1. Inserting test data...")
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
data = insert_sample_prediction(postgres_engine, model_id)
|
|
|
|
print("[TEST] ✓ Data inserted successfully")
|
|
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'REPEAT'
|
|
|
|
print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...")
|
|
workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_1_predict_gate_triggers_continue(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.3.1: Predict Gate Triggers CONTINUE
|
|
|
|
Description:
|
|
Predict response gate determines data should continue despite issues.
|
|
|
|
Expected Behavior:
|
|
- request_predict succeeds
|
|
- mlflow_response_gate for predict returns path_flag='CONTINUE'
|
|
- path_flag_handler calls export workflow with predict data
|
|
- Prediction exported despite quality issues
|
|
|
|
Assertions:
|
|
- Transform and predict completed
|
|
- mlflow_response_gate called for predict
|
|
- Export workflow called with predict data
|
|
- Workflow completes
|
|
"""
|
|
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 = 210"))
|
|
|
|
insert_sql = """
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
(210, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
|
(210, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
print("[TEST] ✓ Data inserted successfully")
|
|
|
|
# Mock mlflow_response_gate for predict to return CONTINUE
|
|
original_mlflow_response_gate = test_activities.mlflow_response_gate
|
|
|
|
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)
|
|
|
|
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!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_2_predict_gate_triggers_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.3.2: Predict Gate Triggers STOP
|
|
|
|
Description:
|
|
Prediction validation fails with STOP policy.
|
|
|
|
Expected Behavior:
|
|
- request_predict succeeds but response invalid
|
|
- mlflow_response_gate for predict returns path_flag='STOP'
|
|
- Workflow exits without export
|
|
|
|
Assertions:
|
|
- Transform completed
|
|
- Predict completed but validation failed
|
|
- Export workflow NOT called
|
|
- Workflow completes without error
|
|
"""
|
|
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 = 205"))
|
|
|
|
insert_sql = """
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
(205, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
|
(205, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
print("[TEST] ✓ Data inserted successfully")
|
|
|
|
# Mock mlflow_response_gate for predict to return STOP
|
|
original_mlflow_response_gate = test_activities.mlflow_response_gate
|
|
|
|
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'],
|
|
}
|
|
|
|
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!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_3_predict_gate_triggers_repeat(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Scenario 2.3.3: Predict Gate Triggers REPEAT
|
|
|
|
Description:
|
|
Predict response gate determines data should repeat last prediction.
|
|
|
|
Expected Behavior:
|
|
- request_predict succeeds but response has issues
|
|
- mlflow_response_gate for predict returns path_flag='REPEAT'
|
|
- path_flag_handler calls repeat_last_prediction activity
|
|
- Last prediction repeated and exported
|
|
|
|
Assertions:
|
|
- Transform and predict completed but validation triggered REPEAT
|
|
- mlflow_response_gate called for predict
|
|
- repeat_last_prediction activity called
|
|
- Export workflow NOT called with current prediction
|
|
- Workflow completes
|
|
"""
|
|
client = temporal_test_env.client
|
|
|
|
print("\n[TEST] 1. Inserting test data and previous prediction...")
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 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))
|
|
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'],
|
|
}
|
|
|
|
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] 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 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!")
|