Files
sientia-dataops-laborious-l…/e2e/test_predictions_batch_main_workflow.py
vitor-aignosi 10c7e292b9 SIENTIAPDE-1646
Refactor ModelMetrics to utilize DriftAnalysis for drift detection

- Replaced ModelAnalysis with DriftAnalysis in the ModelMetrics class to enhance drift detection capabilities.
- Updated method signatures and documentation to reflect the changes in target_name and return values.
- Adjusted data handling to ensure compatibility with the new analysis methods and improved clarity in the drift metrics dataframe preparation.
2026-05-08 16:39:12 -03:00

182 lines
6.1 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 load_scenario_input, 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 sientia_data.laborious_data WHERE model_id = 123'))
insert_sql = """
INSERT INTO sientia_data.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 = load_scenario_input('main_happy_path.json', model_id=123)
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-predictions-batch'),
)
schema_name = 'sientia_data'
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 = load_scenario_input('main_sql_error.json', model_id=128)
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 sientia_data.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 = load_scenario_input('main_missing_required.json', model_id=129)
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 sientia_data.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 sientia_data.laborious_data WHERE model_id = 130'))
conn.execute(
text(
"""
INSERT INTO sientia_data.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 = load_scenario_input('main_invalid_datetime.json', model_id=130)
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 sientia_data.predictions WHERE model_id = 130')
).scalar()
assert count == 0
await handle.terminate('expected failure path in e2e test (invalid datetime column)')