vitor-aignosi
1ce8b9d3a7
Enhance environment configuration and update dependencies - Added new environment variables for PluginStore and MLflow configuration in `.env.example`, including `RUNTIME`, `STORE_BASE_URL`, `STORE_OWNER`, `STORE_REPO`, `STORE_BRANCH`, `STORE_USERNAME`, `STORE_PASSWORD`, `STORE_CACHE_TTL_SECONDS`, `PYPI_SERVER`, `PYPI_USERNAME`, and `PYPI_PASSWORD`. - Updated `git-requirements-mapping.txt` to reflect changes in repository names. - Modified `requirements-light.txt` and `requirements.txt` to upgrade `sientia-dataops-library` to version 1.12.0 and `sientia-mlops-library` to version 0.8.1. - Updated `values.yaml` to include new environment variables for worker runtime and PluginStore configuration. - Refactored E2E tests to utilize new MLflow repository stubs and PluginStore mocks for improved testing accuracy.
386 lines
13 KiB
Python
386 lines
13 KiB
Python
"""
|
|
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import (
|
|
assert_continue,
|
|
assert_repeat,
|
|
assert_stop,
|
|
insert_sample_data,
|
|
make_workflow_id,
|
|
start_and_await_workflow,
|
|
)
|
|
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',
|
|
'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,
|
|
'target': 'sensor_1',
|
|
},
|
|
'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_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_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')
|
|
|
|
|
|
@pytest.fixture
|
|
def bad_data_model(mlflow_repository_stub):
|
|
mlflow_repository_stub.stub_wrapper.transform = MagicMock(
|
|
side_effect=Exception('Bad data model')
|
|
)
|
|
return mlflow_repository_stub.stub_wrapper
|
|
|
|
|
|
@pytest.fixture
|
|
def bad_predict_model(mlflow_repository_stub):
|
|
wrapper = mlflow_repository_stub.stub_wrapper
|
|
|
|
def _good_transform(data):
|
|
result = pd.DataFrame(
|
|
{
|
|
'feature_1': [0.234] * len(data),
|
|
'feature_2': [0.783] * len(data),
|
|
}
|
|
)
|
|
result.index = data.index
|
|
return result, {}
|
|
|
|
wrapper.transform.side_effect = _good_transform
|
|
wrapper.predict = MagicMock(side_effect=Exception('Bad predict model'))
|
|
return wrapper
|
|
|
|
|
|
@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,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
|
|
client = temporal_test_env.client
|
|
model_id = 211
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
|
|
)
|
|
assert_continue(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
|
|
|
|
|
@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,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""Input gate STOP: no export, no MLflow."""
|
|
client = temporal_test_env.client
|
|
model_id = 212
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
|
|
|
|
@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,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""Input gate REPEAT with existing history."""
|
|
client = temporal_test_env.client
|
|
model_id = 213
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
data = insert_sample_prediction(postgres_engine, model_id)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""REPEAT when no prior row in predictions: repeat_last_prediction runs; still no new duplicate export path."""
|
|
client = temporal_test_env.client
|
|
model_id = 214
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-no-history')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
|
|
@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,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 221
|
|
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'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue')
|
|
)
|
|
assert_continue(
|
|
postgres_engine=postgres_engine,
|
|
model_id=model_id,
|
|
prediction_confidence=Decimal(10),
|
|
comments='Unknown MLFlow API error',
|
|
)
|
|
|
|
|
|
@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,
|
|
mlflow_repository_stub,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 222
|
|
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'] = 'STOP'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
|
|
|
|
|
@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,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 223
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
data = insert_sample_prediction(postgres_engine, model_id)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter)."""
|
|
client = temporal_test_env.client
|
|
model_id = 224
|
|
|
|
def all_nan_transform(data):
|
|
result = pd.DataFrame(
|
|
{'feature_1': [np.nan] * len(data), 'feature_2': [np.nan] * len(data)}
|
|
)
|
|
result.index = data.index
|
|
return result, {}
|
|
|
|
mlflow_repository_stub.stub_wrapper.transform = MagicMock(side_effect=all_nan_transform)
|
|
|
|
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': 'STOP', 'CONFIG': {}},
|
|
'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
}
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
|
|
|
|
|
@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,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 231
|
|
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'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue')
|
|
)
|
|
assert_continue(
|
|
postgres_engine=postgres_engine,
|
|
model_id=model_id,
|
|
prediction_confidence=Decimal(10),
|
|
comments='Unknown MLFlow API error',
|
|
)
|
|
|
|
|
|
@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,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 232
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
|
|
@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,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 233
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
data = insert_sample_prediction(postgres_engine, model_id)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
|
input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_4_1_input_empty_data_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""EMPTY_DATA filter with STOP when query returns no rows (offload payload empty)."""
|
|
client = temporal_test_env.client
|
|
model_id = 241
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|