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

SIENTIAPDE-1712

Browse Source 
Enhance E2E testing with MinIO support and update documentation

- Updated `requirements-dev.txt` to include MinIO support in testcontainers.
- Added a new fixture for MinIO container setup in `conftest.py` to facilitate E2E tests involving S3-compatible storage.
- Introduced a new test fixture for activities using a real MinIO container in `conftest.py`.
- Updated E2E test scenarios and documentation to reflect the integration of MinIO for offload uploads and clarified error handling in workflows.
- Refactored existing tests to improve clarity and maintainability.
This commit is contained in:
vitor-aignosi
2026-03-24 09:44:35 -03:00
parent 503d9aa485
commit 0e3ec6463f
11 changed files with 870 additions and 893 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
e2e/conftest.py
View File

@@ -9,6 +9,7 @@ import pandas as pd
import pytest
import pytest_asyncio
from sqlalchemy import create_engine, text
from testcontainers.minio import MinioContainer
from testcontainers.postgres import PostgresContainer
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
@@ -26,6 +27,17 @@ TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'
@pytest_asyncio.fixture(scope='session')
def minio_container():
"""
MinIO S3-compatible storage for E2E tests that exercise real offload uploads.
"""
minio = MinioContainer()
minio.start()
yield minio
minio.stop()
@pytest_asyncio.fixture(scope='session')
def postgres_container():
"""
@@ -187,6 +199,20 @@ def mock_mongo_client():
return mock_client
@pytest.fixture
def notification_inserts(mock_mongo_client):
"""
Mongo insert_one mock used by CoreNotificationHandler for notification persistence.
Yields:
MagicMock for insert_one, reset before each test.
"""
mock_db = mock_mongo_client.__getitem__.return_value
mock_collection = mock_db.__getitem__.return_value
mock_collection.insert_one.reset_mock()
yield mock_collection.insert_one
@pytest_asyncio.fixture
def notification_handler(mock_logger, mock_mongo_client):
"""
@@ -439,6 +465,88 @@ async def test_activities(
await activities.shutdown()
@pytest_asyncio.fixture(scope='function')
async def test_activities_real_minio(
postgres_engine,
postgres_container,
minio_container,
mock_logger,
notification_handler,
metrics_controller,
patch_create_engine,
patch_mlflow,
patch_pi_web_api_repository,
mock_opc_repository,
):
"""
Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests.
"""
minio_client = minio_container.get_client()
if not minio_client.bucket_exists('test-bucket'):
minio_client.make_bucket('test-bucket')
minio_port = minio_container.get_exposed_port(9000)
activities = Activities(
postgres_config={
'host': 'localhost',
'port': postgres_container.get_exposed_port(5432),
'user': 'test',
'password': 'test',
'dbname': 'test',
'min_connections': 1,
'max_connections': 5,
},
mlflow_config={
'host': 'http://localhost',
'port': '5000',
'username': 'test',
'password': 'test',
},
minio_config={
'endpoint_url': f'localhost:{minio_port}',
'access_key': 'minioadmin',
'secret_key': 'minioadmin',
'default_bucket': 'test-bucket',
'retention_hours': 24,
'secure': False,
},
opc_config={},
pi_web_api_config={
'base_url': 'http://localhost:8080',
'auth_type': 'bearer',
'auth_token': 'test_token',
},
logger=mock_logger,
notification_handler=notification_handler,
)
activities.opc_repository = {'1': mock_opc_repository}
try:
yield activities
finally:
await activities.shutdown()
def _worker_activity_list(test_activities: Activities):
return [
test_activities.load_custom_query,
test_activities.load_query_with_minio_offload,
test_activities.cleanup_minio_objects_expired,
test_activities.input_gate,
test_activities.request_transform,
test_activities.mlflow_response_gate,
test_activities.mlflow_content_gate,
test_activities.request_predict,
test_activities.repeat_last_prediction,
test_activities.format_prediction,
test_activities.format_transformed_data,
test_activities.format_default_prediction,
test_activities.write_pi_web_api_data,
test_activities.write_opc_data,
test_activities.export_data_to_postgres,
test_activities.export_payload_to_postgres,
test_activities.write_metrics,
]
@pytest_asyncio.fixture(scope='function')
async def temporal_test_env():
"""Create Temporal test environment."""
@@ -454,24 +562,18 @@ async def temporal_worker(temporal_test_env, test_activities):
temporal_test_env.client,
task_queue='test-queue',
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
activities=[
test_activities.load_custom_query,
test_activities.load_query_with_minio_offload,
test_activities.cleanup_minio_objects_expired,
test_activities.input_gate,
test_activities.request_transform,
test_activities.mlflow_response_gate,
test_activities.mlflow_content_gate,
test_activities.request_predict,
test_activities.repeat_last_prediction,
test_activities.format_prediction,
test_activities.format_transformed_data,
test_activities.format_default_prediction,
test_activities.write_pi_web_api_data,
test_activities.write_opc_data,
test_activities.export_data_to_postgres,
test_activities.export_payload_to_postgres,
test_activities.write_metrics,
],
activities=_worker_activity_list(test_activities),
) as worker:
yield worker
@pytest_asyncio.fixture(scope='function')
async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio):
"""Temporal worker backed by Activities using real MinIO testcontainer."""
async with Worker(
temporal_test_env.client,
task_queue='test-queue',
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
activities=_worker_activity_list(test_activities_real_minio),
) as worker:
yield worker

174
e2e/helpers.py Normal file
View File

@@ -0,0 +1,174 @@
"""
Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
"""
import asyncio
from datetime import datetime
from decimal import Decimal
from typing import Any
from sqlalchemy import text
from sqlalchemy.engine import Engine
async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0):
"""
Start a workflow and wait for its result.
Args:
client: Temporal client from WorkflowEnvironment.
workflow_run: Workflow run method (e.g. PredictionsBatch.run).
input_data: Workflow input payload.
workflow_id: Unique workflow id.
timeout: Max seconds to wait for completion.
Return:
Workflow result value.
"""
handle = await client.start_workflow(
workflow_run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
return await asyncio.wait_for(handle.result(), timeout=timeout)
def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None:
"""
Replace laborious_data rows for a model_id with one row per value (sensor_1..n).
Args:
postgres_engine: SQLAlchemy engine.
model_id: Model id column value.
values: Per-sensor values; use string 'NULL' for SQL NULL.
"""
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
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 assert_prediction(
postgres_engine: Engine,
model_id: int,
prediction: float = 0.5,
prediction_confidence: int | Decimal = 0,
prediction_status: str = 'Good',
comments: str = '',
) -> None:
"""
Assert exactly one prediction row exists for model_id with expected columns.
Args:
postgres_engine: SQLAlchemy engine.
model_id: Expected model_id.
prediction: Expected prediction value.
prediction_confidence: Expected confidence (int or Decimal for numeric column).
prediction_status: Expected status string.
comments: Expected comments string.
"""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
f'ORDER BY created_at ASC'
)
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
row = prediction_rows[0]
assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
f'Expected prediction={prediction}, got {row[1]}'
)
assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
str(prediction_confidence)
), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
def assert_continue(
postgres_engine: Engine,
model_id: int,
prediction_confidence: Decimal = Decimal(2),
comments: str = 'Input data with bad quality',
) -> None:
"""Assert one default-style prediction row after CONTINUE gate path."""
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
f'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'
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: Engine, model_id: int) -> None:
"""Assert no prediction rows for model_id."""
import pytest
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: Engine, model_id: int, last_prediction: tuple) -> None:
"""
Assert two prediction rows for model_id both match last_prediction.
Rows are compared in created_at order for stability.
Args:
postgres_engine: SQLAlchemy engine.
model_id: Model id.
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
"""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(
f'SELECT model_id, prediction, prediction_confidence, prediction_status '
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
f'ORDER BY created_at ASC'
)
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 2, 'Expected two prediction records'
assert prediction_rows[0] == last_prediction, (
f'Expected first row {last_prediction}, got {prediction_rows[0]}'
)
assert prediction_rows[1] == last_prediction, (
f'Expected second row {last_prediction}, got {prediction_rows[1]}'
)
def make_workflow_id(prefix: str) -> str:
"""Build a unique workflow id using a prefix and current timestamp."""
return f'{prefix}-{datetime.now().timestamp()}'

39
e2e/scenarios.md
View File

@@ -2,6 +2,13 @@
This document describes all possible test scenarios for the `predictions_batch` workflow and its child workflows `prediction_process` and `format_and_export_prediction`.
## Running automated E2E tests (`e2e/`)
- **Runtime**: Docker (or a Docker-compatible daemon) must be available so [testcontainers](https://testcontainers.com/) can start **PostgreSQL** and **MinIO** containers.
- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio]`).
- **Invocation**: run only integration-marked tests, for example: `pytest e2e/ -m integration`.
- **MinIO tests**: `e2e/test_minio_offload.py` exercises real S3 uploads; other E2E modules continue to mock MinIO on the worker used by most scenarios.
## Workflow Overview
The `predictions_batch` workflow:
@@ -445,6 +452,8 @@ The `predictions_batch` workflow:
### 3.2 Error Scenarios
These paths do **not** rely on Temporal activity retries for export failures: the write activities run once, errors are handled inside the activity, and the **workflow completes successfully** with degraded metadata on the persisted prediction (`prediction_confidence` and `comments`).
#### Scenario 3.2.1: PI Web API Write Error
**Description**: PI Web API export fails
@@ -453,15 +462,16 @@ The `predictions_batch` workflow:
- PI Web API service unavailable or invalid config
**Expected Behavior**:
- `write_pi_web_api_data` raises exception
- Notification sent
- Workflow fails after retries
- PostgreSQL export may not execute (depends on execution order)
- `write_pi_web_api_data` surfaces the failure (exception handled in the activity layer)
- Notification may be sent
- Workflow **completes** (does not fail)
- Prediction row is still written to PostgreSQL with error confidence **13** and a comment describing the PI error
- Subsequent steps (e.g. OPC, Postgres) still run per workflow order with the updated prediction payload
**Assertions**:
- PI Web API error notification sent
- Workflow fails
- May impact subsequent exports
- PI Web API error notification sent (when applicable)
- Workflow completes
- PostgreSQL contains the prediction with `prediction_confidence` 13 and expected `comments`
---
@@ -473,14 +483,15 @@ The `predictions_batch` workflow:
- OPC server unavailable or invalid configuration
**Expected Behavior**:
- `write_opc_data` raises exception
- Notification sent
- Workflow fails after retries
- `write_opc_data` reports failure without aborting the workflow
- Notification may be sent
- Workflow **completes** (does not fail)
- Prediction row is written to PostgreSQL with OPC error confidence **12** and a comment indicating OPC write issues
**Assertions**:
- OPC error notification sent
- Workflow fails
- PostgreSQL export may not execute
- OPC error notification sent (when applicable)
- Workflow completes
- PostgreSQL contains the prediction with `prediction_confidence` 12 and expected `comments`
---
@@ -497,7 +508,7 @@ The `predictions_batch` workflow:
- `process_pi_web_api_response` detects partial failure
- Error confidence set (13)
- Notification sent for failed tag
- Workflow completes with error confidence
- Workflow completes with error confidence (single activity attempt; no retry loop)
**Assertions**:
- One tag written successfully

74
e2e/test_child_workflows_e2e.py Normal file
View File

@@ -0,0 +1,74 @@
"""
Direct E2E execution of child workflows (smaller surface than PredictionsBatch).
"""
from decimal import Decimal
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.helpers import make_workflow_id, start_and_await_workflow
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
@pytest.mark.asyncio
@pytest.mark.integration
async def test_format_and_export_prediction_default_path_e2e(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
postgres_engine,
):
"""
Run FormatAndExportPrediction with path_flag set (format_default_prediction path).
"""
client = temporal_test_env.client
model_id = 401
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
metadata = {
'metadata': {
'model_id': model_id,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'subworkflow.format_and_export_prediction',
}
}
input_data = {
'metadata': metadata,
'path_flag': 'CONTINUE',
'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'},
'prediction_confidence': 2,
'timestamp': '2024-01-01 12:00:00+00:00',
'model_id': model_id,
'model_name': 'test_model',
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
'comment': 'e2e child workflow default path',
'opc_output_config': {},
'pi_web_api_output_config': {},
'prediction_store_policy': 'lts:1',
}
await start_and_await_workflow(
client,
FormatAndExportPrediction.run,
input_data,
make_workflow_id('e2e-format-export-child'),
)
with postgres_engine.connect() as conn:
row = conn.execute(
text(
f'SELECT prediction, prediction_confidence, prediction_status, comments '
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
)
).fetchone()
assert row is not None
assert row[0] == 0
assert row[1] == Decimal(2)
assert row[2] == 'Bad'
assert row[3] == 'e2e child workflow default path'

124
e2e/test_minio_offload.py Normal file
View File

@@ -0,0 +1,124 @@
"""
E2E-style tests for MinIO offload using a real MinIO testcontainer.
"""
from unittest.mock import patch
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow
from laborious.activities.activities import Activities
from laborious.utils.models import minio_dataframe_payload as mdp
from laborious.workflows.predictions_batch import PredictionsBatch
@pytest.mark.asyncio
@pytest.mark.integration
async def test_load_query_with_minio_offload_writes_object_to_bucket(
postgres_engine,
minio_container,
test_activities_real_minio: Activities,
):
"""
With a tiny offload threshold, query results are uploaded as Parquet to MinIO.
Uses real MinioRepository against testcontainers MinIO (no MinIO mock).
"""
model_id = 501
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
metadata = {
'metadata': {
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': model_id,
'workflow_name': 'predictions_batch',
}
}
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
payload = await test_activities_real_minio.load_query_with_minio_offload(
{
**metadata,
'query': (
'SELECT timestamp, variable, value, created_at '
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
),
'model_name': 'test_model',
'datetime_columns': ['timestamp', 'created_at'],
}
)
assert payload.object_key, 'offloaded payload must reference a MinIO object'
assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'
df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
assert len(df) >= 1
client = minio_container.get_client()
listed = list(client.list_objects('test-bucket', recursive=True))
names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'
@pytest.mark.asyncio
@pytest.mark.integration
async def test_predictions_batch_with_minio_offload_path(
temporal_test_env: WorkflowEnvironment,
temporal_worker_real_minio: Worker,
postgres_engine,
test_activities_real_minio: Activities,
):
"""
Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes.
"""
model_id = 502
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}'))
insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
input_data = {
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': model_id,
'query': (
'SELECT timestamp, variable, value, created_at '
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
),
'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'],
}
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
await start_and_await_workflow(
temporal_test_env.client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-batch-minio-offload'),
)
with postgres_engine.connect() as conn:
count = conn.execute(
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
).scalar()
assert count == 1

290
e2e/test_predictions_batch_format_export.py
View File

@@ -2,18 +2,17 @@
End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
"""
import asyncio
from datetime import datetime
from decimal import Decimal
from typing import Any, cast
from unittest.mock import ANY, AsyncMock, patch, call
from unittest.mock import ANY, AsyncMock, call
import pandas as pd
import pytest
from sientia_do.notifications.models import NotificationLevel
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@@ -56,70 +55,6 @@ def get_base_input_data(model_id):
'query': base_query.format(model_id=model_id),
}
def insert_sample_data(postgres_engine, model_id, values: list):
with postgres_engine.begin() as conn:
conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}"))
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))
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_prediction(
postgres_engine, model_id, prediction: float = 0.5,
prediction_confidence: int = 0, prediction_status: str = 'Good',
comments: str = '',
):
"""
Verify prediction was created with correct values in database
Args:
postgres_engine: Database engine
model_id: Model ID to check
prediction: Expected prediction value (default 0.5 from mock)
prediction_confidence: Expected confidence value (default 0 for normal predictions)
prediction_status: Expected status (default 'Good')
comments: Expected comments (default empty string)
"""
print("\n[TEST] 4. Verifying prediction was created with correct values...")
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, f"Expected one prediction record, got {len(prediction_rows)}"
row = prediction_rows[0]
assert row[0] == model_id, f"Expected model_id={model_id}, got {row[0]}"
assert row[1] == prediction, f"Expected prediction={prediction}, got {row[1]}"
assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}"
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
@pytest.mark.asyncio
@pytest.mark.integration
@@ -130,35 +65,31 @@ async def test_scenario_3_1_1_default_prediction_export(
postgres_engine,
):
"""
Scenario 3.1.1: Default Prediction Export
Scenario 3.1.1: Default prediction export (non-None path_flag).
Description:
Error prediction path creates default prediction.
Triggers input_gate CONTINUE via SPECIFIC_VARIABLES_NULL_VALUES so
PredictionProcess calls FormatAndExportPrediction with path_flag set.
That workflow uses format_default_prediction (not format_prediction) and
skips format_transformed_data / transform Postgres export.
Expected Behavior:
- format_default_prediction called instead of format_prediction
- Default prediction created with error metadata
- Exported to PostgreSQL only
- Transformed data NOT processed
- Metrics written
Assertions:
- format_default_prediction called
- format_prediction NOT called
- format_transformed_data NOT called
- One PostgreSQL export only
- Default values in prediction data
- Comment included
Optional PI Web API and OPC outputs still run when configured.
"""
client = temporal_test_env.client
model_id = 311
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
with postgres_engine.begin() as conn:
conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}"))
conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
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': 'CONTINUE',
'CONFIG': {'variables': ['sensor_1']},
},
}
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
@@ -179,10 +110,9 @@ async def test_scenario_3_1_1_default_prediction_export(
}
}
print("\n[TEST] 2. Starting workflow that should create default prediction...")
workflow_id = f'test-default-prediction-{datetime.now().timestamp()}'
wid = make_workflow_id('test-default-prediction')
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(client, PredictionsBatch.run, input_data, wid)
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
@@ -190,7 +120,7 @@ async def test_scenario_3_1_1_default_prediction_export(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0.5,
'Value': 0,
},
metadata={
'model_id': 311,
@@ -203,7 +133,7 @@ async def test_scenario_3_1_1_default_prediction_export(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
'Value': 2,
},
metadata={
'model_id': 311,
@@ -219,26 +149,48 @@ async def test_scenario_3_1_1_default_prediction_export(
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.assert_has_calls(
[
call('addr_1', 0.5, 'float', ANY,
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
call('addr_2', 0, 'float', ANY,
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
call(
'addr_1',
0,
'float',
ANY,
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
'addr_2',
2,
'float',
ANY,
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
]
)
assert_prediction(postgres_engine, model_id)
with postgres_engine.connect() as conn:
tf_count = conn.execute(
text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
).scalar()
assert tf_count == 0, 'transform export must be skipped when path_flag is set'
assert_prediction(
postgres_engine,
model_id,
prediction=0,
prediction_confidence=Decimal(2),
prediction_status='Bad',
comments='Input data with bad quality',
)
print("\n[TEST] ✓ All assertions passed!")
@@ -274,9 +226,7 @@ async def test_scenario_3_1_2_export_with_opc_only(
model_id = 312
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['opc_output_config'] = {
@@ -295,10 +245,9 @@ async def test_scenario_3_1_2_export_with_opc_only(
}
input_data['pi_web_api_output_config'] = None # No PI Web API config
print("\n[TEST] 2. Starting workflow with OPC only...")
workflow_id = f'test-opc-only-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-only')
)
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.assert_has_calls(
@@ -324,7 +273,6 @@ async def test_scenario_3_1_2_export_with_opc_only(
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@@ -358,9 +306,7 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
model_id = 313
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['pi_web_api_output_config'] = {
@@ -370,10 +316,9 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
}
input_data['opc_output_config'] = None # No OPC config
print("\n[TEST] 2. Starting workflow with PI Web API only...")
workflow_id = f'test-pi-api-only-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-only')
)
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
@@ -412,7 +357,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@@ -445,18 +389,15 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
model_id = 314
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['opc_output_config'] = None # No OPC config
input_data['pi_web_api_output_config'] = None # No PI Web API config
print("\n[TEST] 2. Starting workflow without optional outputs...")
workflow_id = f'test-no-optional-outputs-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-no-optional-outputs')
)
test_activities.pi_web_api_client.write_value.assert_not_called()
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
@@ -464,7 +405,6 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@@ -495,11 +435,9 @@ async def test_scenario_3_1_5_export_without_transformed_data(
model_id = 315
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
with postgres_engine.begin() as conn:
conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
print("[TEST] ✓ Data inserted successfully")
input_data = get_base_input_data(model_id)
input_data['save_transform'] = False # Don't save transformed data
@@ -523,10 +461,9 @@ async def test_scenario_3_1_5_export_without_transformed_data(
}
}
print("\n[TEST] 2. Starting workflow without transformed data export...")
workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-no-transform-export')
)
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
@@ -589,7 +526,6 @@ async def test_scenario_3_1_5_export_without_transformed_data(
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@@ -599,31 +535,20 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
notification_inserts,
):
"""
Scenario 3.2.1: PI Web API Write Error
Description:
PI Web API export fails.
Expected Behavior:
- write_pi_web_api_data raises exception
- Notification sent
- Workflow fails after retries
- PostgreSQL export may not execute (depends on execution order)
Assertions:
- PI Web API error notification sent
- Workflow fails
- May impact subsequent exports
Export failure is handled inside the activity; there is no retry loop. The
workflow completes and PostgreSQL stores prediction_confidence 13 and the
error message in comments.
"""
client = temporal_test_env.client
model_id = 321
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
test_activities.pi_web_api_client.write_value.side_effect = Exception(
"PI Web API service unavailable")
@@ -649,18 +574,17 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
}
}
print("\n[TEST] 2. Starting workflow that should fail on PI Web API write...")
workflow_id = f'test-pi-api-error-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-error')
)
assert_prediction(
postgres_engine, model_id,
prediction_confidence=13,
comments='PI Web API service unavailable',
)
assert notification_inserts.call_count >= 1
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@@ -674,26 +598,14 @@ async def test_scenario_3_2_2_opc_write_error(
"""
Scenario 3.2.2: OPC Write Error
Description:
OPC server write fails.
Expected Behavior:
- write_opc_data raises exception
- Notification sent
- Workflow fails after retries
Assertions:
- OPC error notification sent
- Workflow fails
- PostgreSQL export may not execute
OPC failure is reported without failing the workflow; there is no retry
loop. PostgreSQL stores prediction_confidence 12 and OPC error comments.
"""
client = temporal_test_env.client
model_id = 322
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.return_value = (False, {
@@ -725,10 +637,9 @@ async def test_scenario_3_2_2_opc_write_error(
'confidence_tags': {'tag_2': 'web_id_2'},
}
print("\n[TEST] 2. Starting workflow that should fail on OPC write...")
workflow_id = f'test-opc-error-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error')
)
assert_prediction(
postgres_engine, model_id,
@@ -736,7 +647,6 @@ async def test_scenario_3_2_2_opc_write_error(
comments='Some data could not be written to OPC servers',
)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@@ -750,30 +660,14 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
"""
Scenario 3.2.3: PI Web API Partial Write Error
Description:
Two prediction tags attempt to be written to PI Web API, but only one succeeds.
Expected Behavior:
- write_pi_web_api_data processes response
- process_pi_web_api_response detects partial failure
- Error confidence set (13)
- Notification sent for failed tag
- Workflow completes with error confidence
Assertions:
- One tag written successfully
- One tag failed
- Error confidence set in prediction
- Error notification sent
- Workflow completes
Partial PI write: confidence 13, descriptive comments, workflow completes
without an activity retry loop.
"""
client = temporal_test_env.client
model_id = 323
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
test_activities.pi_web_api_client.write_value = AsyncMock(
side_effect=[
@@ -805,10 +699,9 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
}
}
print("\n[TEST] 2. Starting workflow with partial PI Web API write error...")
workflow_id = f'test-pi-api-partial-error-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error')
)
assert_prediction(
postgres_engine, model_id,
@@ -816,4 +709,3 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.",
)
print("\n[TEST] ✓ All assertions passed!")

283
e2e/test_predictions_batch_main_workflow.py
View File

@@ -3,14 +3,14 @@ End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
"""
import asyncio
from datetime import datetime
import pandas as pd
import pytest
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
@@ -23,55 +23,11 @@ async def test_scenario_1_1_1_happy_path_complete_success(
test_activities: Activities,
postgres_engine,
):
"""
Scenario 1.1.1: Happy Path - Complete Success
Description:
Workflow completes successfully with valid SQL query and all activities succeed.
Process Flow:
1. load_custom_query returns DataFrame with sensor data
2. Workflow prepares prediction input with all configurations
3. prediction_process child workflow executes:
- get_last_timestamp retrieves last processing timestamp
- input_gate validates data quality (passes)
- request_transform calls MLFlow transform (mocked, returns features)
- mlflow_response_gate validates transform response (passes)
- mlflow_content_gate validates transform content (passes)
- request_predict calls MLFlow predict (mocked, returns predictions)
- mlflow_response_gate validates predict response (passes)
- mlflow_content_gate validates predict content (passes)
4. format_and_export_prediction child workflow executes:
- format_prediction formats the prediction data
- format_transformed_data formats transformed data (if save_transform=True)
- export_data_to_postgres saves to database
- write_metrics records execution metrics
Expected Behavior:
- All activities execute successfully without errors
- All gates pass with no quality issues
- Transform and predict operations succeed (mocked)
- Data exported to PostgreSQL predictions table
- Transformed data exported to transformed_data table (if save_transform=True)
- Metrics written successfully
Assertions:
- Workflow completes without raising exceptions
- Data exists in PostgreSQL predictions table with correct model_id
- Data exists in transformed_data table (if save_transform=True)
- Prediction data has expected structure (jsonb with predictions)
- All required fields are populated (model_id, model_name, timestamp, etc)
"""
"""Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms."""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data into PostgreSQL...")
# Insert test data directly into PostgreSQL
# The load_custom_query activity will fetch this data with a real SQL query
with postgres_engine.begin() as conn:
# Clear any existing data for this model_id
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 123"))
# Insert sensor data that the workflow will query
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
@@ -80,9 +36,7 @@ async def test_scenario_1_1_1_happy_path_complete_success(
(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))
print("[TEST] ✓ Data inserted successfully")
# Prepare input data for PredictionsBatch workflow
input_data = {
'metadata': {
'metadata': {
@@ -121,77 +75,44 @@ async def test_scenario_1_1_1_happy_path_complete_success(
'datetime_columns': ['timestamp', 'created_at'],
}
# Start workflow
print("\n[TEST] 2. Starting workflow...")
workflow_id = f'test-predictions-batch-{datetime.now().timestamp()}'
print(f"[TEST] Workflow ID: {workflow_id}")
handle = await client.start_workflow(
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
make_workflow_id('test-predictions-batch'),
)
print("[TEST] ✓ Workflow started")
# Wait for workflow completion with timeout
print("\n[TEST] 3. Waiting for workflow completion (timeout: 60s)...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0) # 60 seconds timeout
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
print("[TEST] ✗ Workflow TIMEOUT after 60 seconds!")
pytest.fail("Workflow execution timed out after 60 seconds")
# Verify data was stored in PostgreSQL - use single connection
schema_name = 'predictions_schema'
predictions_table = 'predictions'
transformed_table = 'transformed_data'
full_predictions_table = f"{schema_name}.{predictions_table}"
full_transformed_table = f"{schema_name}.{transformed_table}"
# Use a single connection for all verification queries
print("\n[TEST] 4. Verifying results in PostgreSQL...")
with postgres_engine.connect() as conn:
# Verify prediction data
result_query = conn.execute(
text(f"SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments FROM {full_predictions_table} WHERE model_id = 123")
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()
print(f"[TEST] Found {len(prediction_rows)} prediction record(s)")
assert len(prediction_rows) == 1, "Expected one prediction record"
# Verify first row has expected structure
assert len(prediction_rows) == 1
row = prediction_rows[0]
print(f"[TEST] Prediction: {row}")
assert row[0] == 123, f"Expected model_id=123, got {row[0]}"
assert row[1] == 0.5, f"Expected prediction=0.5, got {row[1]}"
assert row[2] == 0, f"Expected prediction_confidence=0.9, got {row[2]}"
assert row[3] is not None, f"Expected response_time=0.1, got {row[3]}"
assert row[4] == 'Good', f"Expected prediction_status='Good', got {row[4]}"
assert row[5] == '', f"Expected comments='', got {row[5]}"
print("[TEST] ✓ Prediction data verified")
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] == ''
# Verify transformed data
result_query = conn.execute(
text(f"SELECT model_id, variable, value FROM {full_transformed_table} WHERE model_id = 123")
text(
f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123'
)
)
transformed_rows = result_query.fetchall()
print(f"[TEST] Found {len(transformed_rows)} transformed data record(s)")
assert len(transformed_rows) == 2, "Expected two transformed data records"
row_1 = transformed_rows[0]
print(f"[TEST] Transformed data: {row_1}")
assert row_1[0] == 123, f"Expected model_id=123, got {row_1[0]}"
assert row_1[1] == 'feature_1', f"Expected variable='sensor_1', got {row_1[1]}"
assert float(row_1[2]) == 0.234, f"Expected value=0.234, got {row_1[2]}"
row_2 = transformed_rows[1]
print(f"[TEST] Transformed data: {row_2}")
assert row_2[0] == 123, f"Expected model_id=123, got {row_2[0]}"
assert row_2[1] == 'feature_2', f"Expected variable='sensor_2', got {row_2[1]}"
assert float(row_2[2]) == 0.783, f"Expected value=0.783, got {row_2[2]}"
print("\n[TEST] ✓ All assertions passed!")
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
@@ -202,23 +123,7 @@ async def test_scenario_1_2_1_sql_query_execution_error(
test_activities: Activities,
postgres_engine,
):
"""
Scenario 1.2.1: SQL Query Execution Error
Description:
SQL query fails due to syntax error or connection issue.
Expected Behavior:
- load_custom_query raises exception (caught by Temporal retry policy)
- Notification sent with SQL error details
- After retries, activity may return empty data or workflow may fail
- If empty data returned, workflow completes with early exit via input gate
Assertions:
- Error notification sent
- Workflow completes (either fails or exits early)
- No data in predictions table
"""
"""Invalid SQL: workflow may complete with early exit; no prediction rows."""
client = temporal_test_env.client
input_data = {
@@ -233,7 +138,7 @@ async def test_scenario_1_2_1_sql_query_execution_error(
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 128,
'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =', # Invalid SQL
'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =',
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
@@ -258,34 +163,18 @@ async def test_scenario_1_2_1_sql_query_execution_error(
},
}
print("\n[TEST] 1. Starting workflow with invalid SQL query...")
workflow_id = f'test-sql-error-{datetime.now().timestamp()}'
handle = await client.start_workflow(
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
make_workflow_id('test-sql-error'),
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 2. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed (may have exited early due to empty data)")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
# Verify no predictions were created (regardless of whether workflow failed or exited early)
print("\n[TEST] 3. 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 = 128")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")
count = conn.execute(
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
).scalar()
assert count == 0
@pytest.mark.asyncio
@@ -296,24 +185,9 @@ async def test_scenario_1_2_2_missing_required_parameters(
test_activities: Activities,
postgres_engine,
):
"""
Scenario 1.2.2: Missing Required Parameters
Description:
Essential parameters missing from input.
Expected Behavior:
- Workflow or activity raises KeyError or validation error
- Workflow fails immediately
Assertions:
- Workflow fails with parameter error
- Error notification sent
- No child workflow called
"""
"""Missing query: workflow does not produce predictions and is terminated explicitly."""
client = temporal_test_env.client
# Missing 'query' parameter
input_data = {
'metadata': {
'metadata': {
@@ -326,31 +200,28 @@ async def test_scenario_1_2_2_missing_required_parameters(
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 129,
# 'query' is missing
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
}
print("\n[TEST] 1. Starting workflow with missing required parameter...")
workflow_id = f'test-missing-param-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
id=make_workflow_id('test-missing-param'),
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 2. Waiting for workflow to fail...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
pytest.fail("Expected workflow to fail, but it completed successfully")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
# Let Temporal process a few workflow tasks; for this case, result() can hang.
await asyncio.sleep(2.0)
print("\n[TEST] ✓ All assertions passed!")
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
@@ -361,34 +232,19 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification(
test_activities: Activities,
postgres_engine,
):
"""
Scenario 1.2.3: Invalid Datetime Column Specification
Description:
Datetime column specified doesn't exist in query results.
Expected Behavior:
- load_custom_query may raise KeyError or warning
- Depending on implementation, workflow may fail or continue
- Error notification sent
Assertions:
- Error raised or warning logged
- Workflow behavior depends on error handling policy
"""
"""Invalid datetime column: no predictions persisted; workflow terminated after validation."""
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 = 130"))
insert_sql = """
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')
VALUES (130, 'sensor_1', 23.5, '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': {
@@ -425,26 +281,23 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification(
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['nonexistent_column'], # Column doesn't exist in query result
'datetime_columns': ['nonexistent_column'],
}
print("\n[TEST] 2. Starting workflow with invalid datetime column...")
workflow_id = f'test-invalid-datetime-col-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
id=make_workflow_id('test-invalid-datetime-col'),
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion or failure...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
# Workflow may complete or fail depending on error handling
print("[TEST] ✓ Workflow completed (may have handled error gracefully)")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
# Let Temporal process and surface the failure path internally.
await asyncio.sleep(2.0)
print("\n[TEST] ✓ Test completed!")
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)')

593
e2e/test_predictions_batch_prediction_process.py
View File

@@ -2,56 +2,62 @@
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
"""
import asyncio
from datetime import datetime
from decimal import Decimal
from typing import Any
from unittest.mock import MagicMock, patch
import numpy as np
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 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', # Continue despite issues, not STOP
'CONFIG': {'variables': ['sensor_1']},
},
'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,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['timestamp', 'created_at'],
}
},
'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}'
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 {
@@ -60,97 +66,38 @@ def get_base_input_data(model_id):
'query': base_query.format(model_id=model_id),
}
def insert_sample_data(postgres_engine, model_id, values: list[Any]):
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)
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')
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 = 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]}"
return (model_id, Decimal(10), Decimal(0), 'Good')
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"
@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
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.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
@@ -160,45 +107,19 @@ async def test_scenario_2_1_1_input_gate_triggers_continue(
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
mock_mlflow_models,
):
"""
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
"""
"""Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
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)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
)
assert_continue(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
mock_mlflow_models['transform_model'].predict.assert_not_called()
mock_mlflow_models['predict_model'].predict.assert_not_called()
@pytest.mark.asyncio
@@ -208,45 +129,19 @@ async def test_scenario_2_1_2_input_gate_triggers_stop(
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
mock_mlflow_models,
):
"""
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
"""
"""Input gate STOP: no export, no MLflow."""
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)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
)
assert_stop(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
mock_mlflow_models['transform_model'].predict.assert_not_called()
@pytest.mark.asyncio
@@ -256,57 +151,42 @@ async def test_scenario_2_1_3_input_gate_triggers_repeat(
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
mock_mlflow_models,
):
"""
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
"""
"""Input gate REPEAT with existing history."""
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")
)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat')
)
assert_repeat(postgres_engine, model_id, data)
mock_mlflow_models['transform_model'].predict.assert_not_called()
patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
return model
@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
@@ -318,40 +198,14 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue(
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)
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,
@@ -359,8 +213,6 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue(
comments='Unknown MLFlow API error',
)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
@@ -370,44 +222,18 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(
test_activities: Activities,
postgres_engine,
bad_data_model,
mock_mlflow_models,
):
"""
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)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
)
assert_stop(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
mock_mlflow_models['predict_model'].predict.assert_not_called()
@pytest.mark.asyncio
@@ -419,67 +245,50 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat(
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)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
)
assert_repeat(postgres_engine, model_id, data)
print("\n[TEST] ✓ All assertions passed!")
@pytest.fixture
def bad_predict_model(
patch_mlflow,
mock_mlflow_models
@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,
mock_mlflow_models,
):
model = MagicMock(
predict=MagicMock(
side_effect=Exception("Bad predict model")
)
"""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):
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows})
result.index = data.index
return result
mock_mlflow_models['transform_model'].predict = 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')
)
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
assert_stop(postgres_engine, model_id)
mock_mlflow_models['predict_model'].predict.assert_not_called()
@pytest.mark.asyncio
@@ -491,38 +300,14 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue(
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)
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,
@@ -530,8 +315,6 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue(
comments='Unknown MLFlow API error',
)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
@@ -542,42 +325,16 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop(
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)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop')
)
assert_stop(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
@@ -588,43 +345,35 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat(
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)
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)
print("\n[TEST] ✓ All assertions passed!")
@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)

9
laborious/activities/api.py
View File

@@ -248,8 +248,13 @@ class API(SientiaMonitoring):
metadata=metadata,
)
data['prediction_confidence'] = confidence
data['comments'] = message
# Preserve incoming confidence/comments on successful PI writes.
# Only downgrade confidence or override comments when PI response
# explicitly reports a problem (e.g. partial write mismatch).
if confidence != 0:
data['prediction_confidence'] = confidence
if message:
data['comments'] = message
except Exception as e:
trace = traceback.format_exc()

13
laborious/activities/storage.py
View File

@@ -1,5 +1,3 @@
import json
from temporalio import activity, workflow
from laborious.utils.repository.minio_manager import MinioManager
@@ -195,14 +193,9 @@ class Storage(Postgres, MinioManager):
)
self.error(trace, metadata)
else:
await self.send_notification_async(
metadata=metadata,
notification_id='CLEANUP_MINIO_OBJECTS_EXPIRED',
message='MinIO objects cleaned up successfully',
block='cleanup_minio_objects_expired',
level=NotificationLevel.INFO,
attachment_content=json.dumps(report),
)
# Cleanup success is expected in normal flow; avoid noisy INFO notifications
# that do not impact behavior and can flood observability in test runs.
self.info('MinIO objects cleaned up successfully', metadata)
return report

2
requirements-dev.txt
View File

@@ -13,7 +13,7 @@ types-requests>=2.31.0 # Type stubs for requests
pytest>=7.4.0 # Testing framework
pytest-cov>=4.1.0 # Coverage plugin for pytest
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
testcontainers[postgres] # PostgreSQL containers for E2E tests
testcontainers[postgres,minio] # PostgreSQL and MinIO containers for E2E tests
# Development Tools
ipython>=8.12.0 # Enhanced Python shell