304 lines
10 KiB
Python
304 lines
10 KiB
Python
"""
|
|
End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
import pytest
|
|
|
|
from e2e.helpers import make_workflow_id, start_and_await_workflow
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_1_1_happy_path_complete_success(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms."""
|
|
client = temporal_test_env.client
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123'))
|
|
insert_sql = """
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
(123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
|
(123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
|
(123, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 123,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 123,
|
|
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 123',
|
|
'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'],
|
|
}
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-predictions-batch'),
|
|
)
|
|
|
|
schema_name = 'predictions_schema'
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments '
|
|
f'FROM {schema_name}.predictions WHERE model_id = 123'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1
|
|
row = prediction_rows[0]
|
|
assert row[0] == 123
|
|
assert row[1] == 0.5
|
|
assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}'
|
|
assert row[3] is not None
|
|
assert row[4] == 'Good'
|
|
assert row[5] == ''
|
|
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123'
|
|
)
|
|
)
|
|
transformed_rows = result_query.fetchall()
|
|
assert len(transformed_rows) == 2
|
|
assert transformed_rows[0][0] == 123
|
|
assert transformed_rows[0][1] == 'feature_1'
|
|
assert float(transformed_rows[0][2]) == 0.234
|
|
assert transformed_rows[1][0] == 123
|
|
assert transformed_rows[1][1] == 'feature_2'
|
|
assert float(transformed_rows[1][2]) == 0.783
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_2_1_sql_query_execution_error(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Invalid SQL: workflow may complete with early exit; no prediction rows."""
|
|
client = temporal_test_env.client
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 128,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 128,
|
|
'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =',
|
|
'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',
|
|
},
|
|
}
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-sql-error'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_2_2_missing_required_parameters(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Missing query: workflow does not produce predictions and is terminated explicitly."""
|
|
client = temporal_test_env.client
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 129,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 129,
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
}
|
|
|
|
handle = await client.start_workflow(
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
id=make_workflow_id('test-missing-param'),
|
|
task_queue='test-queue',
|
|
)
|
|
|
|
# Let Temporal process a few workflow tasks; for this case, result() can hang.
|
|
await asyncio.sleep(2.0)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129')
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
await handle.terminate('expected failure path in e2e test (missing required parameters)')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_2_3_invalid_datetime_column_specification(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Invalid datetime column: no predictions persisted; workflow terminated after validation."""
|
|
client = temporal_test_env.client
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130'))
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
"""
|
|
)
|
|
)
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 130,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 130,
|
|
'query': 'SELECT timestamp, variable, value FROM predictions_schema.laborious_data WHERE model_id = 130',
|
|
'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': ['nonexistent_column'],
|
|
}
|
|
|
|
handle = await client.start_workflow(
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
id=make_workflow_id('test-invalid-datetime-col'),
|
|
task_queue='test-queue',
|
|
)
|
|
|
|
# Let Temporal process and surface the failure path internally.
|
|
await asyncio.sleep(2.0)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130')
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
await handle.terminate('expected failure path in e2e test (invalid datetime column)')
|