Snapshot of fix/QTZPOC-13 source tree

Code-only import without upstream history.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
vitor-aignosi
2026-08-19 11:50:01 -03:00
commit f794d9e11e
90 changed files with 24649 additions and 0 deletions

3
e2e/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
"""
End-to-end tests for laborious temporal workflows.
"""

767
e2e/conftest.py Normal file
View File

@@ -0,0 +1,767 @@
"""
Pytest configuration and fixtures for E2E tests.
"""
import sys
from unittest.mock import AsyncMock, MagicMock, patch
# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
_model_analysis_module = MagicMock()
_model_analysis_module.ModelAnalysis = MagicMock
sys.modules.setdefault('sientia', MagicMock())
sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
from io import BytesIO
import joblib
import pandas as pd
import pytest
import pytest_asyncio
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.metrics_controller import MetricsController
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sqlalchemy import create_engine, text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from testcontainers.minio import MinioContainer
from testcontainers.postgres import PostgresContainer
from e2e.opc_test_server import OpcE2ETestServer
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
# Test constants
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'
# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised.
JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model'
JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact'
@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():
"""
Create a PostgreSQL container using testcontainers.
This fixture creates a real PostgreSQL database in a Docker container
that will be used for all tests in the session.
"""
postgres = PostgresContainer('postgres:15')
postgres.start()
yield postgres
postgres.stop()
@pytest_asyncio.fixture
def postgres_engine(postgres_container):
"""
Create SQLAlchemy engine for PostgreSQL test database.
This fixture creates a connection to the PostgreSQL container
created by the postgres_container fixture.
"""
engine = create_engine(postgres_container.get_connection_url())
yield engine
engine.dispose()
def _create_schema_and_tables(engine):
"""
Helper function to create schema and tables in the given engine.
Creates predictions_schema with:
- laborious_data: Input data table for queries
- predictions: Output predictions table
- transformed_data: Output transformed data table
"""
# Use begin() to ensure transaction is properly committed
with engine.begin() as conn:
# Create predictions_schema
conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema'))
# Create laborious_data table (input data from sensors)
create_laborious_data_sql = """
CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
id SERIAL NOT NULL,
model_id int4 NOT NULL,
variable text NOT NULL,
value numeric NULL,
"timestamp" timestamptz NOT NULL,
created_at timestamptz NOT NULL,
PRIMARY KEY (id)
);
"""
conn.execute(text(create_laborious_data_sql))
# Create predictions table
create_predictions_sql = """
CREATE TABLE if not exists predictions_schema.predictions (
id SERIAL NOT NULL ,
model_id int4 NOT NULL,
prediction numeric NULL,
prediction_confidence numeric NOT NULL,
response_time numeric NOT NULL,
prediction_status text NOT NULL,
"timestamp" timestamptz NOT NULL,
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
"comments" text NULL,
PRIMARY KEY (id, created_at)
);
"""
conn.execute(text(create_predictions_sql))
# Create transformed_data table
create_transformed_sql = """
CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
id SERIAL NOT NULL,
model_id int4 NOT NULL,
variable text NOT NULL,
value numeric NULL,
"timestamp" timestamptz NOT NULL,
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
PRIMARY KEY (id)
);
"""
conn.execute(text(create_transformed_sql))
@pytest_asyncio.fixture(autouse=True)
def setup_postgres_schema_and_tables(postgres_engine):
"""
Automatically create necessary schema and tables before each test.
This fixture runs automatically (autouse=True) and ensures
that the predictions_schema and tables exist with the correct structure.
"""
_create_schema_and_tables(postgres_engine)
yield
@pytest_asyncio.fixture
def mock_logger():
"""Mock logger for testing."""
def message(message):
print(f'[LOG] {message}')
def custom_message(message, _metadata=None):
print(f'[LOG] {message}')
logger = MagicMock()
logger.info = MagicMock(side_effect=message)
logger.debug = MagicMock(side_effect=message)
logger.error = MagicMock(side_effect=message)
logger.warning = MagicMock(side_effect=message)
logger.custom_info = MagicMock(side_effect=custom_message)
logger.custom_debug = MagicMock(side_effect=custom_message)
logger.custom_error = MagicMock(side_effect=custom_message)
logger.custom_warning = MagicMock(side_effect=custom_message)
return logger
@pytest_asyncio.fixture
def mock_mongo_client():
"""
Mock MongoDB client to avoid real connections.
This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
allowing us to use a real NotificationHandler instance without connecting to MongoDB.
"""
mock_client = MagicMock()
mock_db = MagicMock()
mock_collection = MagicMock()
# Configure the mock chain: client[database] -> db[collection] -> collection
mock_client.__getitem__.return_value = mock_db
mock_db.__getitem__.return_value = mock_collection
# Mock server_info() to avoid connection attempts
mock_client.server_info = MagicMock()
# Mock insert_one for notifications
mock_collection.insert_one = MagicMock()
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):
"""
Create a real NotificationHandler instance with mocked MongoDB client.
This fixture creates a real CoreNotificationHandler instance but mocks
the underlying MongoDB connection to avoid real database connections.
"""
# Patch MongoClient where it's imported in the handlers module
with patch('sientia_do.notifications.handlers.MongoClient', return_value=mock_mongo_client):
handler = CoreNotificationHandler(
connection_string=TEST_MONGODB_CONNECTION_STRING,
database=TEST_DATABASE_NAME,
logger=mock_logger,
project_name='laborious',
)
yield handler
handler.shutdown()
@pytest_asyncio.fixture
def metrics_controller(mock_logger):
"""Create a real MetricsController instance."""
return MetricsController(logger=mock_logger)
@pytest_asyncio.fixture
def mock_minio_repository():
"""Mock MinIO repository for object storage operations."""
mock_repo = MagicMock()
# Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve()
# can decode the payload if offloading is exercised in an integration scenario.
parquet_df = pd.DataFrame({'a': [1]})
parquet_buffer = BytesIO()
parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
parquet_bytes = parquet_buffer.getvalue()
# sientia_do MinioRepository API
mock_repo.bucket = 'test-bucket'
mock_repo.upload_file = AsyncMock(
side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: {
'minio_object_name': f'sientia/streamlit-connectors/{relative_key}',
'original_filename': relative_key.rsplit('/', 1)[-1],
'uploaded_at': '2024-01-01T00:00:00Z',
'sha256_hash': 'deadbeef',
}
)
mock_repo.download_file = AsyncMock(return_value=parquet_bytes)
mock_repo.list_objects = AsyncMock(return_value=[])
mock_repo.delete_file = AsyncMock()
mock_repo.close = MagicMock()
return mock_repo
@pytest_asyncio.fixture
def mock_pi_web_api_repository():
"""Mock PI Web API repository for PI Web API operations."""
mock_repo = MagicMock()
async def _write_value(web_ids, value, metadata=None, **kwargs):
"""
Mirror successful PI writes: one response item per requested web_id.
write_pi_web_api_data passes the list into process_pi_web_api_response (not a
wrapped {'Items': ...} envelope).
"""
return [{'WebId': wid, 'Errors': []} for wid in web_ids]
mock_repo.write_value = AsyncMock(side_effect=_write_value)
mock_repo.close = MagicMock()
return mock_repo
@pytest_asyncio.fixture
def mock_opc_repository():
"""Mock OPC repository for OPC operations."""
mock_repo = MagicMock()
mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1}))
mock_repo.disconnect = AsyncMock()
return mock_repo
@pytest_asyncio.fixture
async def opc_e2e_server():
"""
In-process asyncua OPC UA server for E2E tests against OpcRepository.
"""
server = OpcE2ETestServer()
await server.start()
try:
yield server
finally:
await server.stop()
@pytest_asyncio.fixture
def patch_create_engine(postgres_engine):
"""Patch create_engine to return test postgres_engine."""
with patch(
'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine
):
yield
@pytest_asyncio.fixture
def patch_minio_repository(mock_minio_repository):
"""Patch MinioRepository to return mock."""
# Patch where Activities resolves the symbol (import binds the original class).
with patch(
'laborious.activities.activities.MinioRepository', return_value=mock_minio_repository
):
yield
@pytest_asyncio.fixture
def patch_pi_web_api_repository(mock_pi_web_api_repository):
"""Patch MLflowRepository to return mock."""
with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
yield
@pytest_asyncio.fixture
def mock_mlflow_models():
"""Create mock models for MLflow load_model methods."""
# Mock transform model - returns DataFrame with same index as input
mock_transform_model = MagicMock()
def mock_transform_predict(data):
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
print(data.to_csv())
print(data.index)
result = pd.DataFrame(
{
'feature_1': [0.234] * num_rows,
'feature_2': [0.783] * num_rows,
}
)
result.index = data.index
return result
mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
# Mock predict model - returns array/list of predictions
mock_predict_model = MagicMock()
def mock_predict_predict(data):
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
return [0.5] * num_rows
mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
# Mock PyFuncModel for compressed models
mock_pyfunc_model = MagicMock()
mock_pyfunc_model._model_impl = MagicMock()
mock_pyfunc_model._model_impl.python_model = mock_transform_model
return {
'transform_model': mock_transform_model,
'predict_model': mock_predict_model,
'pyfunc_model': mock_pyfunc_model,
}
def _build_joblib_fallback_estimator():
"""
Build a real sklearn estimator for the joblib fallback artifact.
`request_predict` appends a string `timestamp` column to the transformed
data, so the estimator selects the transform output columns explicitly and
drops the rest.
"""
features = ['feature_1', 'feature_2']
estimator = Pipeline(
steps=[
(
'select',
ColumnTransformer([('features', 'passthrough', features)], remainder='drop'),
),
('regressor', LinearRegression()),
]
)
train = pd.DataFrame(
{
'feature_1': [0.0, 1.0, 2.0, 3.0],
'feature_2': [0.0, 1.0, 0.0, 1.0],
'timestamp': ['2024-01-01 12:00:00'] * 4,
}
)
estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3)
return estimator
@pytest.fixture
def joblib_fallback_artifact(tmp_path):
"""
Real artifact directory holding a joblib-compressed sklearn estimator.
Mirrors what the joblib fallback finds after downloading `prediction_model`:
a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an
`MLmodel` placeholder.
Return:
dict with 'path' (the artifact directory) and 'estimator' (the object
that was serialized, so tests can compute the expected prediction).
"""
artifact_dir = tmp_path / 'prediction_model'
artifact_dir.mkdir()
estimator = _build_joblib_fallback_estimator()
joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3)
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
return {'path': str(artifact_dir), 'estimator': estimator}
@pytest.fixture
def joblib_fallback_artifact_without_pickle(tmp_path):
"""Artifact directory with no top-level `.pkl`, so the joblib fallback also fails."""
artifact_dir = tmp_path / 'prediction_model_without_pickle'
artifact_dir.mkdir()
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
return str(artifact_dir)
@pytest_asyncio.fixture
def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
"""Patch mlflow module in repository with load_model mocks."""
mock_mlflow = MagicMock()
# Mock sklearn.load_model
def mock_sklearn_load_model(model_uri):
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
if 'data_model' in model_uri or 'transform' in model_uri.lower():
return mock_mlflow_models['transform_model']
return mock_mlflow_models['predict_model']
mock_mlflow.sklearn = MagicMock()
mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
# Mock pyfunc.load_model
def mock_pyfunc_load_model(model_uri):
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
if 'artifacts' in model_uri or 'tmp' in model_uri:
return mock_mlflow_models['pyfunc_model']
if 'data_model' in model_uri or 'transform' in model_uri.lower():
return mock_mlflow_models['transform_model']
return mock_mlflow_models['predict_model']
mock_mlflow.pyfunc = MagicMock()
mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
# Mock pytorch.load_model
mock_mlflow.pytorch = MagicMock()
mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
# Mock other mlflow methods that might be called
mock_mlflow.set_tracking_uri = MagicMock()
mock_mlflow.get_run = MagicMock(
return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))
)
mock_mlflow.tracking = MagicMock()
mock_mlflow_client = MagicMock(
search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
search_model_versions=MagicMock(
return_value=[
MagicMock(
current_stage='Production', version='1', source='runs:/artifacts/test_run_id'
)
]
),
)
# Only the joblib fallback calls download_artifacts directly on the client.
mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path'])
mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
yield mock_mlflow
@pytest.fixture
def mlflow_client(patch_mlflow):
"""The MlflowClient mock the repository holds as `self.client`."""
return patch_mlflow.tracking.MlflowClient.return_value
@pytest_asyncio.fixture(scope='function')
async def test_activities(
postgres_engine,
postgres_container,
mock_logger,
notification_handler,
metrics_controller,
mock_minio_repository,
patch_create_engine,
patch_minio_repository,
patch_mlflow,
patch_pi_web_api_repository,
mock_opc_repository,
):
"""
Create Activities instance with test dependencies.
This fixture creates a real Activities instance with:
- PostgreSQL database (via testcontainers)
- Mocked MinIO client
- Real NotificationHandler and MetricsController (with mocked underlying services)
"""
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={
# Host:port only; Minio() prepends http(s):// from the secure flag.
'endpoint_url': 'localhost:9000',
'access_key': 'test',
'secret_key': 'test',
'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:
# Cleanup - ALWAYS runs, even if test fails
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."""
env = await WorkflowEnvironment.start_time_skipping()
async with env:
yield env
@pytest_asyncio.fixture(scope='function')
async def temporal_worker(temporal_test_env, test_activities):
"""Create Temporal worker with test activities."""
async with Worker(
temporal_test_env.client,
task_queue='test-queue',
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
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
@pytest_asyncio.fixture(scope='function')
async def test_activities_real_opc(
postgres_engine,
postgres_container,
opc_e2e_server: OpcE2ETestServer,
mock_logger,
notification_handler,
metrics_controller,
patch_create_engine,
patch_minio_repository,
patch_mlflow,
patch_pi_web_api_repository,
):
"""
Activities with a real OpcRepository connected to the in-process OPC UA server.
"""
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': 'localhost:9000',
'access_key': 'test',
'secret_key': 'test',
'default_bucket': 'test-bucket',
'retention_hours': 24,
'secure': False,
},
opc_config={
'1': {
'id': '1',
'server_name': 'e2e-opc',
'url': opc_e2e_server.url,
'server_uri': opc_e2e_server.url,
'cert_path': None,
'private_key_path': None,
'server_cert_path': None,
'reconnection_interval': 0,
}
},
pi_web_api_config={
'base_url': 'http://localhost:8080',
'auth_type': 'bearer',
'auth_token': 'test_token',
},
logger=mock_logger,
notification_handler=notification_handler,
)
await activities.init_opc()
repo = activities.opc_repository['1']
assert repo._session_ready.is_set(), 'OPC E2E server connection failed during init_opc'
try:
yield activities
finally:
await activities.shutdown()
@pytest_asyncio.fixture(scope='function')
async def temporal_worker_real_opc(temporal_test_env, test_activities_real_opc):
"""Temporal worker backed by Activities using the in-process OPC UA server."""
async with Worker(
temporal_test_env.client,
task_queue='test-queue',
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
activities=_worker_activity_list(test_activities_real_opc),
) as worker:
yield worker

186
e2e/helpers.py Normal file
View File

@@ -0,0 +1,186 @@
"""
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 = None,
comments_contains: str | None = None,
) -> 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 exact comments string (optional).
comments_contains: Substring expected in comments when queued (optional).
"""
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]}"
)
if comments is not None:
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
if comments_contains is not None:
assert comments_contains in row[4], (
f"Expected comments to contain '{comments_contains}', 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."""
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.
"""
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()}'

192
e2e/opc_test_server.py Normal file
View File

@@ -0,0 +1,192 @@
"""
In-process OPC UA server for E2E tests (asyncua).
Provides writable prediction/confidence nodes and optional write faults
(Tier-1 BadSessionIdInvalid via PreWrite callback).
"""
from __future__ import annotations
import socket
from dataclasses import dataclass
from typing import TYPE_CHECKING
from asyncua import Server, ua
from asyncua.common.callback import CallbackType
from asyncua.common.utils import ServiceError
if TYPE_CHECKING:
from asyncua.common.node import Node
UNKNOWN_NODE_ID = 'ns=99;i=9999'
@dataclass(frozen=True)
class OpcE2ENodeIds:
"""NodeId strings used in opc_output_config for E2E workflows."""
prediction: str
confidence: str
unknown: str = UNKNOWN_NODE_ID
class OpcE2ETestServer:
"""
Ephemeral asyncua server with Laborious E2E variables and controllable faults.
Args:
host: Bind address (default 127.0.0.1).
"""
def __init__(self, host: str = '127.0.0.1') -> None:
self._host = host
self._server: Server | None = None
self._prediction_node: Node | None = None
self._confidence_node: Node | None = None
self._session_bad_on_write = False
self._url: str | None = None
self._node_ids: OpcE2ENodeIds | None = None
@property
def url(self) -> str:
if self._url is None:
raise RuntimeError('OPC E2E server is not started')
return self._url
@property
def node_ids(self) -> OpcE2ENodeIds:
if self._node_ids is None:
raise RuntimeError('OPC E2E server is not started')
return self._node_ids
def set_session_bad_on_write(self, enabled: bool) -> None:
"""
When enabled, every client Write is rejected with BadSessionIdInvalid.
Args:
enabled (bool): Turn Tier-1 session fault injection on or off.
"""
self._session_bad_on_write = enabled
async def start(self) -> OpcE2ENodeIds:
"""
Start the OPC UA server on a free TCP port.
Return:
OpcE2ENodeIds: NodeId strings for prediction and confidence tags.
"""
port = _free_port(self._host)
self._url = f'opc.tcp://{self._host}:{port}/freeopcua/server/'
server = Server()
server.set_endpoint(self._url)
await server.init()
server.iserver.callback_service.addListener(
CallbackType.PreWrite,
self._pre_write_callback,
)
idx = await server.register_namespace('http://sientia.test/laborious-e2e')
e2e_object = await server.nodes.objects.add_object(idx, 'LaboriousE2E')
prediction = await e2e_object.add_variable(
idx,
'Prediction',
ua.Variant(0.0, ua.VariantType.Float),
)
confidence = await e2e_object.add_variable(
idx,
'Confidence',
ua.Variant(0.0, ua.VariantType.Float),
)
await prediction.set_writable()
await confidence.set_writable()
await server.start()
self._server = server
self._prediction_node = prediction
self._confidence_node = confidence
self._node_ids = OpcE2ENodeIds(
prediction=prediction.nodeid.to_string(),
confidence=confidence.nodeid.to_string(),
)
return self._node_ids
async def stop(self) -> None:
"""Stop the OPC UA server and release the listening port."""
if self._server is not None:
await self._server.stop()
self._server = None
self._prediction_node = None
self._confidence_node = None
self._url = None
self._node_ids = None
self._session_bad_on_write = False
async def read_prediction(self) -> float:
"""
Read the current prediction variable value from the address space.
Return:
float: Stored prediction value.
"""
if self._prediction_node is None:
raise RuntimeError('OPC E2E server is not started')
value = await self._prediction_node.read_value()
return float(value)
async def read_confidence(self) -> float:
"""
Read the current confidence variable value from the address space.
Return:
float: Stored confidence value.
"""
if self._confidence_node is None:
raise RuntimeError('OPC E2E server is not started')
value = await self._confidence_node.read_value()
return float(value)
async def _pre_write_callback(self, event, _service) -> None:
# Only client writes are faulted. asyncua's own clock task writes CurrentTime
# through this same PreWrite dispatch, and raising there kills its time_task;
# the stored exception then resurfaces as a teardown failure in server.stop().
if self._session_bad_on_write and getattr(event, 'is_external', False):
raise ServiceError(ua.StatusCodes.BadSessionIdInvalid)
def _free_port(host: str) -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind((host, 0))
return int(sock.getsockname()[1])
def build_opc_output_config(
node_ids: OpcE2ENodeIds,
*,
prediction_tag: str | None = None,
confidence_tag: str | None = None,
prediction_only: bool = False,
server_key: str = '1',
) -> dict[str, dict]:
"""
Build opc_output_config for PredictionsBatch using real server NodeIds.
Args:
node_ids (OpcE2ENodeIds): Node ids from OpcE2ETestServer.
prediction_tag (str | None): Override prediction NodeId (default: node_ids.prediction).
confidence_tag (str | None): Override confidence NodeId (default: node_ids.confidence).
prediction_only (bool): When True, omit confidence_tags (single write per activity).
server_key (str): OPC server id key in opc_output_config.
Return:
dict: opc_output_config payload for workflow input.
"""
pred = prediction_tag if prediction_tag is not None else node_ids.prediction
conf = confidence_tag if confidence_tag is not None else node_ids.confidence
server_config: dict = {
'prediction_tags': {pred: {'data_type': 'float'}},
}
if not prediction_only:
server_config['confidence_tags'] = {conf: {'data_type': 'float'}}
return {server_key: server_config}

564
e2e/scenarios.md Normal file
View File

@@ -0,0 +1,564 @@
# Test Scenarios for Predictions Batch Workflow
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.
- **OPC tests (real server)**: `e2e/test_opc_real_server.py` uses an in-process **asyncua** server and real `OpcRepository` (`test_activities_real_opc`). Scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 are covered there. Other E2E modules keep the OPC mock.
- Run only OPC real-server tests: `pytest e2e/test_opc_real_server.py -m "integration and opc"`.
## Workflow Overview
The `predictions_batch` workflow:
1. Loads data using a custom SQL query
2. Prepares prediction configuration
3. Delegates to `prediction_process` child workflow which:
- Retrieves last timestamp for incremental processing
- Applies input data quality gates
- Executes MLFlow transform operation
- Validates transform response
- Executes MLFlow predict operation
- Validates predict response
- Delegates to `format_and_export_prediction` child workflow
4. The `format_and_export_prediction` workflow:
- Formats prediction data (normal or default)
- Exports to PI Web API (optional)
- Exports to OPC server (optional)
- Exports to PostgreSQL
- Writes metrics
---
## 1. Predictions Batch - Main Workflow Scenarios
### 1.1 Success Scenarios
#### Scenario 1.1.1: Happy Path - Complete Success
**Description**: Workflow completes successfully with valid SQL query and all activities succeed
**Input**:
- Valid `schedule_name`, `model_name`, `model_id`
- Valid `query` returning non-empty DataFrame
- Valid `schema`, `table_name`, `transform_table_name`
- Optional `datetime_columns` for timestamp parsing
- Optional `input_filters`, `mlflow_transform_filters`, `mlflow_predict_filters`
- Optional `path_priority`, `opc_output_config`, `pi_web_api_output_config`
**Expected Behavior**:
- `load_custom_query` returns DataFrame with data
- Workflow prepares prediction input with all configurations
- `prediction_process` child workflow executes successfully
- All gates pass with no issues
- Transform and predict operations succeed
- Data exported to PostgreSQL
- Metrics written
**Assertions**:
- SQL query executed once
- `prediction_process` workflow called with correct parameters
- Data exists in PostgreSQL (predictions table)
- Metrics recorded
- No errors raised
---
### 1.2 Error Scenarios
#### Scenario 1.2.1: SQL Query Execution Error
**Description**: SQL query fails due to syntax error or connection issue
**Input**:
- Invalid SQL query (syntax error)
- Or database connection unavailable
**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
---
#### Scenario 1.2.2: Missing Required Parameters
**Description**: Essential parameters missing from input
**Input**:
- Missing `query` or `model_id` or `schema` or `table_name`
**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
---
#### Scenario 1.2.3: Invalid Datetime Column Specification
**Description**: Datetime column specified doesn't exist in query results
**Input**:
- `datetime_columns: ['nonexistent_column']`
- Query results don't have this column
**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
---
## 2. Prediction Process - Child Workflow Scenarios
### 2.1 Input gate Early Exit Scenarios
#### Scenario 2.1.1: Input Gate Triggers CONTINUE
**Description**: Input gate determines data should use previous prediction
**Input**:
- Data that should continue with input data as prediction
- `input_filters` configured with `POLICY: 'CONTINUE'`
- `path_priority` includes CONTINUE
**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
#### Scenario 2.1.2: Input Gate Triggers STOP
**Description**: Input data quality gate fails with STOP policy
**Input**:
- Data with EMPTY_DATA or other critical issues
- `input_filters` configured with `POLICY: 'STOP'`
**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
#### Scenario 2.1.3: Input Gate Triggers REPEAT
**Description**: Input gate determines data should repeat last prediction
**Input**:
- Data with quality issues that require using previous prediction
- `input_filters` configured with `POLICY: 'REPEAT'`
- `path_priority` includes REPEAT
**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
---
### 2.2 Transform gate Early Exit Scenarios
#### Scenario 2.2.1: Transform Gate Triggers CONTINUE
**Description**: Transform response gate determines data should continue despite issues
**Input**:
- Valid input data
- Transform response has quality issues but policy is CONTINUE
- `mlflow_transform_filters` configured with `POLICY: 'CONTINUE'`
- `path_priority` includes CONTINUE
**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
---
#### Scenario 2.2.2: Transform Gate Triggers STOP
**Description**: Transform response validation fails with STOP policy
**Input**:
- Valid input data
- Transform response has critical errors
- `mlflow_transform_filters` configured with `POLICY: 'STOP'`
**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
---
#### Scenario 2.2.3: Transform Gate Triggers REPEAT
**Description**: Transform response gate determines data should repeat last prediction
**Input**:
- Valid input data
- Transform response has quality issues that require using previous prediction
- `mlflow_transform_filters` configured with `POLICY: 'REPEAT'`
- `path_priority` includes REPEAT
**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
---
### 2.3 Predict gate Early Exit Scenarios
#### Scenario 2.3.1: Predict Gate Triggers CONTINUE
**Description**: Predict response gate determines data should continue despite issues
**Input**:
- Valid input and transform data
- Predict response has quality issues but policy is CONTINUE
- `mlflow_predict_filters` configured with `POLICY: 'CONTINUE'`
- `path_priority` includes CONTINUE
**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
---
#### Scenario 2.3.2: Predict Gate Triggers STOP
**Description**: Prediction validation fails with STOP policy
**Input**:
- Valid input and transform
- Predict response has critical errors
- `mlflow_predict_filters` configured with `POLICY: 'STOP'`
**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
---
#### Scenario 2.3.3: Predict Gate Triggers REPEAT
**Description**: Predict response gate determines data should repeat last prediction
**Input**:
- Valid input and transform data
- Predict response has quality issues that require using previous prediction
- `mlflow_predict_filters` configured with `POLICY: 'REPEAT'`
- `path_priority` includes REPEAT
**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
---
## 3. Format and Export Prediction - Child Workflow Scenarios
### 3.1 Success Scenarios
#### Scenario 3.1.1: Default Prediction Export
**Description**: Error prediction path creates default prediction
**Input**:
- `path_flag: 'ERROR'` or other non-None value (not STOP/CONTINUE/REPEAT)
- `comment` provided with error details
**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
---
#### Scenario 3.1.2: Export with OPC only
**Description**: Export to PostgreSQL and OPC server only (no PI Web API)
**Input**:
- `path_flag: None`
- `opc_output_config` configured with valid OPC settings
- `pi_web_api_output_config: None` or `{}`
**Expected Behavior**:
- Normal formatting
- PostgreSQL export executed
- OPC export executed
- PI Web API activity skipped
- Metrics written with OPC metrics
**Assertions**:
- PI Web API activity NOT called
- OPC activity called
- PostgreSQL export called
- Metrics written with `opc_metrics` populated
---
#### Scenario 3.1.3: Export with PI Web API only
**Description**: Export to PostgreSQL and PI Web API only (no OPC)
**Input**:
- `path_flag: None`
- `pi_web_api_output_config` configured with valid PI Web API settings
- `opc_output_config: None` or `{}`
**Expected Behavior**:
- Normal formatting
- PostgreSQL export executed
- PI Web API export executed
- OPC activity skipped
- Metrics written without OPC metrics
**Assertions**:
- OPC activity NOT called
- PI Web API activity called
- PostgreSQL export called
- Metrics written with empty `opc_metrics`
---
#### Scenario 3.1.4: Export Without Optional Outputs
**Description**: Export only to PostgreSQL (no OPC or PI Web API)
**Input**:
- `path_flag: None`
- `opc_output_config: None` or `{}`
- `pi_web_api_output_config: None` or `{}`
**Expected Behavior**:
- Normal formatting
- Only PostgreSQL export executed
- OPC and PI Web API activities skipped
- Metrics written without OPC metrics
**Assertions**:
- PI Web API activity NOT called
- OPC activity NOT called
- PostgreSQL export called
- Metrics written with empty `opc_metrics`
---
#### Scenario 3.1.5: Export Without Transformed Data
**Description**: Only prediction exported, no transform table
**Input**:
- `path_flag: None`
- `transformed_data: None` or `save_transform: False`
- `opc_output_config: None` or `{}`
- `pi_web_api_output_config: None` or `{}`
**Expected Behavior**:
- Only prediction formatted and exported
- Transform export skipped
- Single PostgreSQL write
**Assertions**:
- `format_transformed_data` NOT called
- One PostgreSQL export
- Transform table remains empty
---
### 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
**Input**:
- Valid prediction
- PI Web API service unavailable or invalid config
**Expected Behavior**:
- `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 (when applicable)
- Workflow completes
- PostgreSQL contains the prediction with `prediction_confidence` 13 and expected `comments`
---
#### Scenario 3.2.2: OPC Write Error
**Description**: OPC server write fails
**Input**:
- Valid prediction
- OPC server unavailable or invalid configuration
**Expected Behavior**:
- `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 (when applicable)
- Workflow completes
- PostgreSQL contains the prediction with `prediction_confidence` 12 and expected `comments`
---
#### Scenario 3.2.4: OPC Session / Channel Bad* (Tier-1)
**Description**: OPC write fails with a Tier-1 session or channel status (e.g. `BadSessionIdInvalid`) while transport may still appear open on the client
**Input**:
- Valid prediction and OPC output config
- Mock or server returning Tier-1 `UaStatusCodeError` on write (no write retry in the same activity)
**Expected Behavior**:
- `write_opc_data` fails forward for affected tags; background reconnect may be scheduled if `OPC_RECONNECTION_INTERVAL` allows
- Workflow **completes**
- PostgreSQL row uses **`prediction_confidence` 14** and comment prefix `OPC UA session/channel error:` (including OPC status name)
- `opc_write_attempts_total` records `result=BadSessionIdInvalid` (or matching status); no second write attempt in the same activity
**Assertions**:
- Workflow completes
- `prediction_confidence = 14`
- `comments` matches `OPC UA session/channel error:%`
- Generic OPC error confidence **12** is not used for this case
**Reference**: [docs/opc-communication.md](../docs/opc-communication.md), plan `.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md`
---
#### Scenario 3.2.5: OPC Write Blocked During Reconnect
**Description**: A write is attempted while the repository is reconnecting (session not ready)
**Input**:
- Valid prediction
- Simulated slow reconnect (e.g. delayed `connect`) or concurrent writes where the first triggers reconnect
**Expected Behavior**:
- Second write (or parallel write) is rejected **immediately** when reconnect is in progress or `_session_ready` is cleared — **without** calling `write_value`
- No wait/sleep on the write path; no duplicate `connect` from parallel writers (connection lock)
- `prediction_confidence = 14`, `comments = OPC UA reconnect in progress` (distinguish from Tier-1 `Bad*` via comment prefix in SQL)
**Assertions**:
- At most one reconnect sequence (`disconnect` + `connect`) for the overlapping window
- No write retry after failure
- Tests in `test_opc_repository` (unit) and optional e2e in `test_predictions_batch_format_export.py`
---
#### 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
**Input**:
- Valid prediction
- Two prediction tags configured
- PI Web API returns partial success (one tag succeeds, one fails)
**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 (single activity attempt; no retry loop)
**Assertions**:
- One tag written successfully
- One tag failed
- Error confidence set in prediction
- Error notification sent
- Workflow completes
---

View File

@@ -0,0 +1,76 @@
"""
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'

136
e2e/test_minio_offload.py Normal file
View File

@@ -0,0 +1,136 @@
"""
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

211
e2e/test_opc_real_server.py Normal file
View File

@@ -0,0 +1,211 @@
"""
E2E tests for OPC export using an in-process asyncua server and real OpcRepository.
Covers scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 from e2e/scenarios.md.
Mock-based OPC tests remain in test_predictions_batch_format_export.py.
"""
import asyncio
import pytest
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 e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config
from e2e.test_predictions_batch_format_export import get_base_input_data
from laborious.activities.activities import Activities
from laborious.activities.opc import OPC_RECONNECT_IN_PROGRESS_COMMENT
from laborious.utils.repository.opc_repository import OpcRepository
from laborious.workflows.predictions_batch import PredictionsBatch
async def _reconnect_under_held_lock(
repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event
) -> None:
"""
Hold the connection lock until released, then reconnect.
The lock must stay held for the whole workflow, and how long the workflow takes to
reach the OPC write is not bounded — a fixed sleep here lets `_reconnect_locked`
finish first, which sets `_session_ready` again and the write then succeeds.
Args:
repo (OpcRepository): Connected repository.
acquired (asyncio.Event): Set once the lock is held, before any write can race.
release (asyncio.Event): Awaited while holding the lock; set by the caller.
"""
async with repo._connection_lock:
acquired.set()
await release.wait()
await repo._reconnect_locked()
@pytest.mark.asyncio
@pytest.mark.integration
@pytest.mark.opc
async def test_scenario_3_1_2_export_with_opc_only_real_server(
temporal_test_env: WorkflowEnvironment,
temporal_worker_real_opc: Worker,
test_activities_real_opc: Activities,
opc_e2e_server: OpcE2ETestServer,
postgres_engine,
):
"""
Scenario 3.1.2 (real OPC): connect, write prediction and confidence, verify server values.
"""
client = temporal_test_env.client
model_id = 412
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
input_data['pi_web_api_output_config'] = None
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-opc-real-happy'),
)
test_activities_real_opc.pi_web_api_client.write_value.assert_not_called()
assert await opc_e2e_server.read_prediction() == pytest.approx(0.5)
assert await opc_e2e_server.read_confidence() == pytest.approx(0.0)
assert_prediction(postgres_engine, model_id)
@pytest.mark.asyncio
@pytest.mark.integration
@pytest.mark.opc
async def test_scenario_3_2_2_opc_write_error_real_server(
temporal_test_env: WorkflowEnvironment,
temporal_worker_real_opc: Worker,
test_activities_real_opc: Activities,
opc_e2e_server: OpcE2ETestServer,
postgres_engine,
):
"""
Scenario 3.2.2 (real OPC): unknown NodeId yields generic write failure (confidence 12).
"""
client = temporal_test_env.client
model_id = 422
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
node_ids = opc_e2e_server.node_ids
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = build_opc_output_config(
node_ids,
prediction_tag=UNKNOWN_NODE_ID,
confidence_tag=UNKNOWN_NODE_ID,
)
input_data['pi_web_api_output_config'] = None
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-opc-real-bad-node'),
)
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=12,
comments='Some data could not be written to OPC servers',
)
@pytest.mark.asyncio
@pytest.mark.integration
@pytest.mark.opc
async def test_scenario_3_2_4_opc_session_bad_real_server(
temporal_test_env: WorkflowEnvironment,
temporal_worker_real_opc: Worker,
test_activities_real_opc: Activities,
opc_e2e_server: OpcE2ETestServer,
postgres_engine,
):
"""
Scenario 3.2.4 (real OPC): server PreWrite fault injects BadSessionIdInvalid (confidence 14).
"""
client = temporal_test_env.client
model_id = 424
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
opc_e2e_server.set_session_bad_on_write(True)
try:
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = build_opc_output_config(
opc_e2e_server.node_ids,
prediction_only=True,
)
input_data['pi_web_api_output_config'] = None
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-opc-real-session-bad'),
)
finally:
opc_e2e_server.set_session_bad_on_write(False)
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=14,
comments_contains='OPC UA session/channel error: BadSessionIdInvalid',
)
@pytest.mark.asyncio
@pytest.mark.integration
@pytest.mark.opc
async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
temporal_test_env: WorkflowEnvironment,
temporal_worker_real_opc: Worker,
test_activities_real_opc: Activities,
opc_e2e_server: OpcE2ETestServer,
postgres_engine,
):
"""
Scenario 3.2.5 (real OPC): writes rejected while reconnect holds the connection lock.
"""
client = temporal_test_env.client
model_id = 425
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
repo = test_activities_real_opc.opc_repository['1']
repo._session_ready.clear()
acquired = asyncio.Event()
release = asyncio.Event()
reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release))
await acquired.wait()
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
input_data['pi_web_api_output_config'] = None
try:
await start_and_await_workflow(
client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-opc-real-reconnect-block'),
)
finally:
release.set()
await reconnect_task
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=14,
comments_contains=OPC_RECONNECT_IN_PROGRESS_COMMENT,
)

View File

@@ -0,0 +1,800 @@
"""
End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
"""
from decimal import Decimal
from typing import Any, cast
from unittest.mock import AsyncMock, call
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
base_input_data = {
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 301,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 301',
'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'],
}
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),
}
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_1_default_prediction_export(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.1: Default prediction export (non-None path_flag).
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.
Optional PI Web API and OPC outputs still run when configured.
"""
client = temporal_test_env.client
model_id = 311
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'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
wid = make_workflow_id('test-default-prediction')
await start_and_await_workflow(client, PredictionsBatch.run, input_data, wid)
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
call(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
},
metadata={
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 2,
},
metadata={
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
],
any_order=True,
)
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.assert_has_calls(
[
call(
'addr_1',
0,
'float',
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
'addr_2',
2,
'float',
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
]
)
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',
)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_2_export_with_opc_only(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.2: Export with OPC only
Description:
Export to PostgreSQL and OPC server only (no PI Web API).
Expected Behavior:
- Normal formatting
- PostgreSQL export executed
- OPC export executed
- PI Web API activity skipped
- Metrics written with OPC metrics
Assertions:
- PI Web API activity NOT called
- OPC activity called
- PostgreSQL export called
- Metrics written with opc_metrics populated
"""
client = temporal_test_env.client
model_id = 312
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
input_data['pi_web_api_output_config'] = None # No PI Web API config
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(
[
call(
'addr_1',
0.5,
'float',
{
'model_id': 312,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
'addr_2',
0,
'float',
{
'model_id': 312,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
]
)
test_activities.pi_web_api_client.write_value.assert_not_called()
assert_prediction(postgres_engine, model_id)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_3_export_with_pi_web_api_only(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.3: Export with PI Web API only
Description:
Export to PostgreSQL and PI Web API only (no OPC).
Expected Behavior:
- Normal formatting
- PostgreSQL export executed
- PI Web API export executed
- OPC activity skipped
- Metrics written without OPC metrics
Assertions:
- OPC activity NOT called
- PI Web API activity called
- PostgreSQL export called
- Metrics written with empty opc_metrics
"""
client = temporal_test_env.client
model_id = 313
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
input_data = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = None # No OPC config
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(
[
call(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0.5,
},
metadata={
'model_id': 313,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
},
metadata={
'model_id': 313,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
],
any_order=True,
)
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.assert_not_called()
assert_prediction(postgres_engine, model_id)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_4_export_without_optional_outputs(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.4: Export Without Optional Outputs
Description:
Export only to PostgreSQL (no OPC or PI Web API).
Expected Behavior:
- Normal formatting
- Only PostgreSQL export executed
- OPC and PI Web API activities skipped
- Metrics written without OPC metrics
Assertions:
- PI Web API activity NOT called
- OPC activity NOT called
- PostgreSQL export called
- Metrics written with empty opc_metrics
"""
client = temporal_test_env.client
model_id = 314
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
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
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)
opc_write_data.assert_not_called()
assert_prediction(postgres_engine, model_id)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_5_export_without_transformed_data(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.5: Export Without Transformed Data
Description:
Only prediction exported, no transform table.
Expected Behavior:
- Only prediction formatted and exported
- Transform export skipped
- Single PostgreSQL write
Assertions:
- format_transformed_data NOT called
- One PostgreSQL export
- Transform table remains empty
"""
client = temporal_test_env.client
model_id = 315
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}')
)
input_data = get_base_input_data(model_id)
input_data['save_transform'] = False # Don't save transformed data
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
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(
[
call(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0.5,
},
metadata={
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
},
metadata={
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
],
any_order=True,
)
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.assert_has_calls(
[
call(
'addr_1',
0.5,
'float',
{
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
'addr_2',
0,
'float',
{
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
]
)
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(
f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
)
)
count = result_query.scalar()
assert count == 0, f'Expected transform table to be empty, but found {count} records'
assert_prediction(postgres_engine, model_id)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_1_pi_web_api_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
notification_inserts,
):
"""
Scenario 3.2.1: PI Web API Write Error
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
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
test_activities.pi_web_api_client.write_value.side_effect = Exception(
'PI Web API service unavailable'
)
input_data = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
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
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_2_opc_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.2: OPC Write Error
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
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.return_value = (
False,
{
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
'message': 'OPC server unavailable',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'OPC server unavailable',
},
)
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error')
)
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=12,
comments='Some data could not be written to OPC servers',
)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_4_opc_session_bad_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.4: OPC session/channel Tier-1 Bad* (e.g. BadSessionIdInvalid).
PostgreSQL stores prediction_confidence 14 and a stable session error comment.
"""
client = temporal_test_env.client
model_id = 324
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
opc_write_data.return_value = (
False,
{
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
'message': 'BadSessionIdInvalid',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'BadSessionIdInvalid',
'opc_error_kind': 'session_bad',
'opc_status': 'BadSessionIdInvalid',
},
)
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
input_data['pi_web_api_output_config'] = None
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-session-bad')
)
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=14,
comments='OPC UA session/channel error: BadSessionIdInvalid',
)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_3_pi_web_api_partial_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.3: PI Web API Partial Write Error
Partial PI write: confidence 13, descriptive comments, workflow completes
without an activity retry loop.
"""
client = temporal_test_env.client
model_id = 323
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
test_activities.pi_web_api_client.write_value = AsyncMock(
side_effect=[
# Prediction batch: two web_ids requested, only one acknowledged.
[{'WebId': 'web_id_1', 'Errors': []}],
# Confidence write succeeds.
[{'WebId': 'web_id_2', 'Errors': []}],
]
)
input_data = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1', 'tag_3': 'web_id_3'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error')
)
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=13,
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.",
)

View File

@@ -0,0 +1,186 @@
"""
End-to-end tests for the joblib fallback on the MLflow model load paths.
`joblib.load` is NOT mocked here: the artifact fixture holds a real sklearn
estimator serialized with `joblib.dump(..., compress=3)`, so a successful run
proves the fallback really deserializes a compressed pickle and serves it.
"""
import pandas as pd
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.conftest import JOBLIB_FALLBACK_FLAVOR_ERROR, JOBLIB_FALLBACK_MODEL_NAME
from e2e.helpers import (
assert_prediction,
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_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
def build_input_data(model_id: int, model_name: str) -> dict:
"""Workflow input for a plain transform + predict run of a given model."""
return {
'schedule_name': 'test-schedule',
'model_name': model_name,
'model_id': model_id,
'query': base_query.format(model_id=model_id),
'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'],
}
def expected_fallback_prediction(estimator) -> float:
"""Prediction the fixture estimator produces for the mocked transform output."""
features = pd.DataFrame(
{
'feature_1': [0.234],
'feature_2': [0.783],
'timestamp': ['2024-01-01 12:00:00'],
}
)
return float(estimator.predict(features)[0])
def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
"""Assert exactly one Good prediction row whose value matches `expected`."""
with postgres_engine.connect() as conn:
rows = conn.execute(
text(
'SELECT prediction, prediction_status '
'FROM predictions_schema.predictions '
f'WHERE model_id = {model_id}'
)
).fetchall()
assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
assert abs(float(rows[0][0]) - expected) < tolerance, (
f'Expected prediction close to {expected}, got {rows[0][0]}'
)
assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"
def logged_messages(mock_logger, level: str) -> str:
"""Concatenate every message logged at a given custom level."""
calls = getattr(mock_logger, f'custom_{level}').call_args_list
return ' '.join(str(logged.args[0]) for logged in calls)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_recovered_by_joblib_fallback(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
joblib_fallback_artifact,
mlflow_client,
mock_logger,
):
"""Flavor loader fails, joblib fallback loads the real compressed pickle, prediction lands."""
client = temporal_test_env.client
model_id = 271
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback')
)
assert_single_prediction_close(
postgres_engine,
model_id,
expected_fallback_prediction(joblib_fallback_artifact['estimator']),
)
assert mlflow_client.download_artifacts.call_count == 1
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged_messages(mock_logger, 'warning')
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_without_fallback_is_unchanged(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
mlflow_client,
mock_mlflow_models,
):
"""A model whose flavor loader works keeps the current result and downloads nothing."""
client = temporal_test_env.client
model_id = 272
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, 'test_model')
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-no-fallback')
)
assert_prediction(postgres_engine, model_id, prediction=0.5)
mock_mlflow_models['predict_model'].predict.assert_called()
mlflow_client.download_artifacts.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_reports_both_causes_when_fallback_fails(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
joblib_fallback_artifact_without_pickle,
mlflow_client,
mock_logger,
):
"""Both attempts fail: workflow stops and the logged traceback carries both causes."""
mlflow_client.download_artifacts.return_value = joblib_fallback_artifact_without_pickle
client = temporal_test_env.client
model_id = 273
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback-failed')
)
assert_stop(postgres_engine, model_id)
logged = logged_messages(mock_logger, 'debug')
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged
assert 'No .pkl file found in prediction_model artifact' in logged

View File

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

View File

@@ -0,0 +1,385 @@
"""
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
"""
from decimal import Decimal
from unittest.mock import MagicMock
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,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['timestamp', 'created_at'],
}
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
def get_base_input_data(model_id):
return {
**base_input_data,
'model_id': model_id,
'query': base_query.format(model_id=model_id),
}
def insert_sample_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(patch_mlflow):
model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model')))
patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
return model
@pytest.fixture
def bad_predict_model(patch_mlflow, mock_mlflow_models):
model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model')))
def mock_sklearn_load_model(model_uri):
if 'data_model' in model_uri or 'transform' in model_uri.lower():
return mock_mlflow_models['transform_model']
return model
patch_mlflow.sklearn = MagicMock()
patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
return model
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_1_input_gate_triggers_continue(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
mock_mlflow_models,
):
"""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)
mock_mlflow_models['transform_model'].predict.assert_not_called()
mock_mlflow_models['predict_model'].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,
mock_mlflow_models,
):
"""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)
mock_mlflow_models['transform_model'].predict.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,
mock_mlflow_models,
):
"""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)
mock_mlflow_models['transform_model'].predict.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,
mock_mlflow_models,
):
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)
mock_mlflow_models['predict_model'].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,
mock_mlflow_models,
):
"""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')
)
assert_stop(postgres_engine, model_id)
mock_mlflow_models['predict_model'].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)