vitor-aignosi
6b9bb38968
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.
630 lines
21 KiB
Python
630 lines
21 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,
|
|
bad_data_model,
|
|
):
|
|
"""
|
|
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.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(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_predict_model,
|
|
):
|
|
"""
|
|
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
|
|
|
|
model_id = 231
|
|
|
|
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_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()}'
|
|
|
|
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',
|
|
)
|
|
|
|
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,
|
|
bad_predict_model,
|
|
):
|
|
"""
|
|
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
|
|
|
|
model_id = 232
|
|
|
|
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['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()}'
|
|
|
|
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_3_3_predict_gate_triggers_repeat(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_predict_model,
|
|
):
|
|
"""
|
|
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
|
|
|
|
model_id = 233
|
|
|
|
print("\n[TEST] 1. Inserting test data and previous 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")
|
|
|
|
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()}'
|
|
|
|
await start_and_await_workflow(client, input_data, workflow_id)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
print("\n[TEST] ✓ All assertions passed!")
|