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

SIENTIAPDE-1478

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2026-01-15 16:17:47 -03:00
parent b5b86db10c
commit 0cd6ae660a
1 changed files with 43 additions and 246 deletions
Download Patch File Download Diff File Expand all files Collapse all files

283
e2e/test_predictions_batch_prediction_process.py
View File

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