SIENTIAPDE-1478

Update coverage source in pyproject.toml, add testcontainers for PostgreSQL in requirements-dev.txt, increment image tag and adjust probe delays in values.yaml, and refine condition checks in format_and_export_prediction.py and mlflow.py. Additionally, enhance test coverage in test_gates.py.
This commit is contained in:
vitor-aignosi
2026-01-14 16:01:03 -03:00
parent 34e8cce282
commit 63017754a8
13 changed files with 3001 additions and 8 deletions

3
e2e/__init__.py Normal file
View File

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

415
e2e/conftest.py Normal file
View File

@@ -0,0 +1,415 @@
"""
Pytest configuration and fixtures for E2E tests.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pandas as pd
import pytest
import pytest_asyncio
from sqlalchemy import create_engine, text
from testcontainers.postgres import PostgresContainer
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
# Test constants
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'
@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={}):
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_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()
# Mock repository methods
mock_repo.put_parquet_from_dataframe = AsyncMock(return_value='test-object-key')
mock_repo.get_parquet_as_dataframe = AsyncMock(return_value=pd.DataFrame())
mock_repo.minio_bucket = 'test-bucket'
return mock_repo
@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."""
with patch('laborious.utils.repository.minio_repository.MinioRepository', return_value=mock_minio_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,
}
@pytest_asyncio.fixture
def patch_mlflow(mock_mlflow_models):
"""Patch mlflow module in repository with load_model mocks."""
mock_mlflow = MagicMock()
# Mock sklearn.load_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 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 '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.tracking.MlflowClient = MagicMock(return_value=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'
)])
))
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
yield mock_mlflow
@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,
):
"""
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={
'endpoint_url': 'http://localhost:9000',
'access_key': 'test',
'secret_key': 'test',
'region_name': 'us-east-1',
'default_bucket': 'test-bucket',
},
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,
)
try:
yield activities
finally:
# Cleanup - ALWAYS runs, even if test fails
await activities.shutdown()
@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=[
test_activities.load_custom_query,
test_activities.get_last_timestamp,
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.write_metrics,
],
) as worker:
yield worker

561
e2e/scenarios.md Normal file
View File

@@ -0,0 +1,561 @@
# 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`.
## 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
---
### 2.4 Error Scenarios
#### Scenario 2.4.1: MLFlow Transform API Error
**Description**: MLFlow transform request fails
**Input**:
- Valid input data
- MLFlow service unavailable or returns error
**Expected Behavior**:
- `request_transform` raises exception
- Notification sent with MLFlow error details
- Workflow fails after retry attempts
**Assertions**:
- Exception raised from transform activity
- Error notification sent
- Workflow fails
- Export NOT called
---
#### Scenario 2.4.2: MLFlow Predict API Error
**Description**: MLFlow predict request fails
**Input**:
- Valid input and transform data
- MLFlow predict service unavailable
**Expected Behavior**:
- `request_predict` raises exception
- Notification sent
- Workflow fails after retries
**Assertions**:
- Transform succeeded
- Predict raised exception
- Error notification sent
- Workflow fails
---
## 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: 'STOP'` or other non-None value
- `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 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.3: Export Without Transformed Data
**Description**: Only prediction exported, no transform table
**Input**:
- `path_flag: None`
- `transformed_data: None`
**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
#### Scenario 3.2.1: PostgreSQL Export Error - Predictions Table
**Description**: Failed to write predictions to database
**Input**:
- Valid formatted prediction
- PostgreSQL connection fails or table doesn't exist
**Expected Behavior**:
- `export_data_to_postgres` raises exception
- Notification sent with database error
- Workflow fails after retries
**Assertions**:
- Exception raised from export activity
- Error notification sent
- Workflow fails
- Metrics NOT written (activity doesn't execute)
---
#### Scenario 3.2.2: 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` raises exception
- Notification sent
- Workflow fails after retries
- PostgreSQL export may not execute (depends on execution order)
**Assertions**:
- PI Web API error notification sent
- Workflow fails
- May impact subsequent exports
---
#### Scenario 3.2.3: OPC Write Error
**Description**: OPC server write fails
**Input**:
- Valid prediction
- OPC server unavailable or invalid configuration
**Expected Behavior**:
- `write_opc_data` raises exception
- Notification sent
- Workflow fails after retries
**Assertions**:
- OPC error notification sent
- Workflow fails
- PostgreSQL export may not execute
---
## 4. End-to-End Integration Scenarios
### 4.1 Complete Success Path
#### Scenario 4.1.1: Full Pipeline Success with All Features
**Description**: Complete workflow execution with all optional features enabled
**Input**:
- Valid SQL query returning data
- All configurations provided (OPC, PI Web API, filters, policies)
- MLFlow services available
- All databases available
**Expected Behavior**:
- SQL query loads data
- Input gate passes
- MLFlow transform succeeds
- MLFlow predict succeeds
- All validations pass
- Prediction formatted
- Transformed data formatted
- Both exported to PostgreSQL
- PI Web API write succeeds
- OPC write succeeds
- Metrics written
**Assertions**:
- All activities executed in correct order
- All three workflows execute (batch, process, export)
- All exports succeed
- All tables have data
- All external systems updated
- Metrics recorded
---
### 4.2 Error Recovery Integration
#### Scenario 4.2.1: Transform Error with Repeat Fallback
**Description**: Transform fails, workflow repeats last prediction
**Input**:
- Valid input
- MLFlow transform fails
- REPEAT policy configured
- Previous prediction exists
**Expected Behavior**:
- Transform fails
- Filter detects error
- Path handler triggers REPEAT
- Last prediction retrieved and re-exported
- Workflow completes successfully
**Assertions**:
- Transform attempted
- Error handled gracefully
- Last prediction copied
- Workflow completes without exception
---

View File

@@ -0,0 +1,710 @@
"""
End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
"""
import asyncio
from datetime import datetime
import pandas as pd
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@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
Description:
Error prediction path creates default prediction.
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
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 301"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(301, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(301, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock mlflow_response_gate for predict to return a non-None path_flag
# Any non-None path_flag that's not STOP/CONTINUE/REPEAT will be passed to format_and_export_prediction
# which will then call format_default_prediction
from unittest.mock import patch
original_mlflow_response_gate = test_activities.mlflow_response_gate
async def mock_mlflow_response_gate(input_data):
# Only return error path_flag for predict, not transform
if input_data.get('type') == 'predict':
metadata = input_data.get('metadata', {})
# Return a path_flag that will be passed to format_and_export_prediction
# but won't trigger early exit (not STOP/CONTINUE/REPEAT)
# The path_flag_handler only returns True for STOP/CONTINUE/REPEAT
# So any other value will make it return False and continue to export
return 'ERROR', -1, 'Error: Prediction validation failed'
# For transform, return normal (None)
return await original_mlflow_response_gate(input_data)
with patch.object(test_activities, 'mlflow_response_gate', side_effect=mock_mlflow_response_gate):
input_data = {
'metadata': {
'metadata': {
'model_id': 301,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'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': 'CONTINUE', 'config': {}}, # CONTINUE allows workflow to proceed
},
'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'],
}
print("\n[TEST] 2. Starting workflow that should create default prediction...")
workflow_id = f'test-default-prediction-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying default prediction was created...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT prediction, prediction_confidence, comments FROM predictions_schema.predictions WHERE model_id = 301")
)
prediction_rows = result_query.fetchall()
# Should have a default prediction with error comment
assert len(prediction_rows) >= 1, "Expected at least one default prediction"
if len(prediction_rows) > 0:
row = prediction_rows[0]
# Default predictions typically have specific characteristics
# The exact values depend on format_default_prediction implementation
print(f"[TEST] Default prediction found: {row}")
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_2_export_without_optional_outputs(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.2: 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
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 302"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(302, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(302, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
input_data = {
'metadata': {
'metadata': {
'model_id': 302,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 302,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 302',
'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': None, # No OPC config
'pi_web_api_output_config': None, # No PI Web API 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'],
}
print("\n[TEST] 2. Starting workflow without optional outputs...")
workflow_id = f'test-no-optional-outputs-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying only PostgreSQL export was executed...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT model_id FROM predictions_schema.predictions WHERE model_id = 302")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_1_3_export_without_transformed_data(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.1.3: 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
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 303"))
conn.execute(text("DELETE FROM predictions_schema.transformed_data WHERE model_id = 303"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(303, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(303, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
input_data = {
'metadata': {
'metadata': {
'model_id': 303,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 303,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 303',
'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': False, # Don't save transformed data
'prediction_store_policy': 'lts:1',
'model_config': {
'retention_minutes': 0,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['timestamp', 'created_at'],
}
print("\n[TEST] 2. Starting workflow without transformed data export...")
workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying only prediction was exported...")
with postgres_engine.connect() as conn:
# Verify prediction exists
result_query = conn.execute(
text("SELECT model_id FROM predictions_schema.predictions WHERE model_id = 303")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record"
# Verify transformed data table is empty
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = 303")
)
count = result_query.scalar()
assert count == 0, f"Expected transform table to be empty, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_1_postgres_export_error_predictions_table(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.1: PostgreSQL Export Error - Predictions Table
Description:
Failed to write predictions to database.
Expected Behavior:
- export_data_to_postgres raises exception
- Notification sent with database error
- Workflow fails after retries
Assertions:
- Exception raised from export activity
- Error notification sent
- Workflow fails
- Metrics NOT written (activity doesn't execute)
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 304"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(304, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(304, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock export_data_to_postgres to raise an exception
from unittest.mock import patch
original_export = test_activities.export_data_to_postgres
call_count = {'count': 0}
async def mock_export_data_to_postgres(*args, **kwargs):
call_count['count'] += 1
# Only fail on predictions table export, not transform table
if call_count['count'] == 1: # First call is predictions table
raise Exception("PostgreSQL connection failed")
return await original_export(*args, **kwargs)
with patch.object(test_activities, 'export_data_to_postgres', side_effect=mock_export_data_to_postgres):
input_data = {
'metadata': {
'metadata': {
'model_id': 304,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 304,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 304',
'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'],
}
print("\n[TEST] 2. Starting workflow that should fail on PostgreSQL export...")
workflow_id = f'test-postgres-error-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow to fail...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
pytest.fail("Expected workflow to fail, but it completed successfully")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
# Verify no predictions were created
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 304")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_2_pi_web_api_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.2: PI Web API Write Error
Description:
PI Web API export fails.
Expected Behavior:
- write_pi_web_api_data raises exception
- Notification sent
- Workflow fails after retries
- PostgreSQL export may not execute (depends on execution order)
Assertions:
- PI Web API error notification sent
- Workflow fails
- May impact subsequent exports
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 305"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(305, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(305, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock write_pi_web_api_data to raise an exception
from unittest.mock import patch
def mock_write_pi_web_api_data(*args, **kwargs):
raise Exception("PI Web API service unavailable")
with patch.object(test_activities, 'write_pi_web_api_data', side_effect=mock_write_pi_web_api_data):
input_data = {
'metadata': {
'metadata': {
'model_id': 305,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 305,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 305',
'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': {
'endpoint': 'test_endpoint',
'prediction_tags': {'prediction': 'test_pred_tag'},
'confidence_tags': {'confidence': 'test_conf_tag'},
},
'save_transform': True,
'prediction_store_policy': 'lts:1',
'model_config': {
'retention_minutes': 0,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['timestamp', 'created_at'],
}
print("\n[TEST] 2. Starting workflow that should fail on PI Web API write...")
workflow_id = f'test-pi-api-error-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow to fail...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
pytest.fail("Expected workflow to fail, but it completed successfully")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_3_opc_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.3: OPC Write Error
Description:
OPC server write fails.
Expected Behavior:
- write_opc_data raises exception
- Notification sent
- Workflow fails after retries
Assertions:
- OPC error notification sent
- Workflow fails
- PostgreSQL export may not execute
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 306"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(306, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(306, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock write_opc_data to raise an exception
from unittest.mock import patch
def mock_write_opc_data(*args, **kwargs):
raise Exception("OPC server unavailable")
with patch.object(test_activities, 'write_opc_data', side_effect=mock_write_opc_data):
input_data = {
'metadata': {
'metadata': {
'model_id': 306,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 306,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 306',
'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': {
'server_name': 'test_server',
'tags': {'prediction': 'test_tag'},
},
'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'],
}
print("\n[TEST] 2. Starting workflow that should fail on OPC write...")
workflow_id = f'test-opc-error-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow to fail...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
pytest.fail("Expected workflow to fail, but it completed successfully")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
print("\n[TEST] ✓ All assertions passed!")

View File

@@ -0,0 +1,278 @@
"""
End-to-end tests for PredictionsBatch workflow - Integration scenarios.
"""
import asyncio
from datetime import datetime
from unittest.mock import patch
import pandas as pd
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_4_1_1_full_pipeline_success_with_all_features(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 4.1.1: Full Pipeline Success with All Features
Description:
Complete workflow execution with all optional features enabled.
Expected Behavior:
- SQL query loads data
- Input gate passes
- MLFlow transform succeeds
- MLFlow predict succeeds
- All validations pass
- Prediction formatted
- Transformed data formatted
- Both exported to PostgreSQL
- PI Web API write succeeds (mocked)
- OPC write succeeds (mocked)
- Metrics written
Assertions:
- All activities executed in correct order
- All three workflows execute (batch, process, export)
- All exports succeed
- All tables have data
- All external systems updated (mocked)
- Metrics recorded
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 401"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(401, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(401, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(401, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
input_data = {
'metadata': {
'metadata': {
'model_id': 401,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 401,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 401',
'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': {
'server_name': 'test_server',
'tags': {'prediction': 'test_tag'},
},
'pi_web_api_output_config': {
'endpoint': 'test_endpoint',
'prediction_tags': {'prediction': 'test_pred_tag'},
'confidence_tags': {'confidence': 'test_conf_tag'},
},
'save_transform': True,
'prediction_store_policy': 'lts:1',
'model_config': {
'retention_minutes': 0,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['timestamp', 'created_at'],
}
print("\n[TEST] 2. Starting full pipeline workflow...")
workflow_id = f'test-full-pipeline-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying all exports and data...")
with postgres_engine.connect() as conn:
# Verify prediction data
result_query = conn.execute(
text("SELECT model_id, prediction, prediction_confidence FROM predictions_schema.predictions WHERE model_id = 401")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record"
# Verify transformed data
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = 401")
)
count = result_query.scalar()
assert count == 2, f"Expected two transformed data records, but found {count}"
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_4_2_1_transform_error_with_repeat_fallback(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 4.2.1: Transform Error with Repeat Fallback
Description:
Transform fails, workflow repeats last prediction.
Expected Behavior:
- Transform fails
- Filter detects error
- Path handler triggers REPEAT
- Last prediction retrieved and re-exported
- Workflow completes successfully
Assertions:
- Transform attempted
- Error handled gracefully
- Last prediction copied
- Workflow completes without exception
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data and previous prediction...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 403"))
conn.execute(text("DELETE FROM predictions_schema.predictions WHERE model_id = 403"))
# Insert input data
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(403, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(403, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
# Insert a previous prediction to repeat
insert_prediction_sql = """
INSERT INTO predictions_schema.predictions
(model_id, prediction, prediction_confidence, response_time, prediction_status, timestamp, created_at, comments)
VALUES
(403, 0.85, 95, 0.15, 'Good', '2024-01-01 11:00:00+00:00', '2024-01-01 11:00:00+00:00', 'Previous successful prediction')
"""
conn.execute(text(insert_prediction_sql))
print("[TEST] ✓ Data and previous prediction inserted")
# Mock request_transform to return an error response
def mock_request_transform(*args, **kwargs):
return {
'success': False, # This will trigger API_ERROR filter
'content': pd.DataFrame(),
}
with patch.object(test_activities, 'request_transform', side_effect=mock_request_transform):
input_data = {
'metadata': {
'metadata': {
'model_id': 403,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 403,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 403',
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
'input_filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
},
'mlflow_transform_filters': {
'API_ERROR': {'policy': 'REPEAT', 'config': {}}, # REPEAT on error
},
'mlflow_predict_filters': {
'API_ERROR': {'policy': 'STOP', 'config': {}},
},
'path_priority': ['REPEAT', 'STOP', 'CONTINUE'], # REPEAT first
'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'],
}
print("\n[TEST] 2. Starting workflow that should trigger REPEAT...")
workflow_id = f'test-repeat-fallback-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying last prediction was repeated...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 403")
)
count = result_query.scalar()
# Should have at least 2 predictions (original + repeated)
assert count >= 1, f"Expected at least one prediction (repeated), but found {count} records"
print("\n[TEST] ✓ All assertions passed!")

View File

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

View File

@@ -0,0 +1,573 @@
"""
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
"""
import asyncio
from datetime import datetime
from unittest.mock import patch
import pandas as pd
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@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,
):
"""
Scenario 2.1.1: Input Gate Triggers CONTINUE
Description:
Input gate determines data should use previous prediction.
Expected Behavior:
- input_gate returns path_flag='CONTINUE'
- path_flag_handler calls export workflow with input data directly
- MLFlow transform and predict skipped
- Data exported as-is
Assertions:
- input_gate called
- MLFlow operations NOT called
- Export workflow called with original data
- Workflow completes
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 201"))
# Insert data with some null values (quality issue)
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(201, 'sensor_1', NULL, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(201, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
input_data = {
'metadata': {
'metadata': {
'model_id': 201,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 201,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
'input_filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'policy': 'CONTINUE', # Continue despite issues, not STOP
'config': {'variables': ['sensor_1']},
},
},
'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'],
}
print("\n[TEST] 2. Starting workflow with CONTINUE policy...")
workflow_id = f'test-continue-policy-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed successfully")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying prediction was created despite warnings...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = 201")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record despite warnings"
# Assert prediction value is 0 and other fields
row = prediction_rows[0]
assert row[1] == 0, f"Expected prediction=0, got {row[1]}"
assert row[2] == 2, f"Expected prediction_confidence=0, got {row[2]}"
assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
assert row[4] == 'Input data with bad quality', f"Expected comments='Input data with bad quality', got {row[4]}"
print("\n[TEST] ✓ All assertions passed!")
@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,
):
"""
Scenario 2.1.2: Input Gate Triggers STOP
Description:
Input data quality gate fails with STOP policy.
Expected Behavior:
- input_gate returns path_flag='STOP'
- path_flag_handler detects STOP
- Workflow returns early without calling MLFlow
- No prediction exported
Assertions:
- input_gate called
- path_flag_handler returns True (early exit)
- MLFlow transform NOT called
- Export workflow NOT called
- Workflow completes without error
"""
client = temporal_test_env.client
print("\n[TEST] 1. Ensuring no data exists (empty data will trigger STOP)...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 202"))
print("[TEST] ✓ Data cleared")
input_data = {
'metadata': {
'metadata': {
'model_id': 202,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 202,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 202',
'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'],
}
print("\n[TEST] 2. Starting workflow that should stop at input gate...")
workflow_id = f'test-input-stop-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion (should exit early)...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed (exited early as expected)")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying no predictions were created...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 202")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")
@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,
):
"""
Scenario 2.3.2: Predict Gate Triggers STOP
Description:
Prediction validation fails with STOP policy.
Expected Behavior:
- request_predict succeeds but response invalid
- mlflow_response_gate for predict returns path_flag='STOP'
- Workflow exits without export
Assertions:
- Transform completed
- Predict completed but validation failed
- Export workflow NOT called
- Workflow completes without error
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 205"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(205, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(205, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock request_predict to return an error response
def mock_request_predict(*args, **kwargs):
return {
'success': False, # This will trigger API_ERROR filter
'content': [],
}
with patch.object(test_activities, 'request_predict', side_effect=mock_request_predict):
input_data = {
'metadata': {
'metadata': {
'model_id': 205,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 205,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 205',
'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': {}}, # STOP on predict error
},
'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'],
}
print("\n[TEST] 2. Starting workflow that should stop at predict gate...")
workflow_id = f'test-predict-stop-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow completion (should exit early)...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
print("[TEST] ✓ Workflow completed (exited early as expected)")
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
print("\n[TEST] 4. Verifying no predictions were created...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 205")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_2_1_mlflow_transform_api_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 2.4.1: MLFlow Transform API Error
Description:
MLFlow transform request fails.
Expected Behavior:
- request_transform raises exception
- Notification sent with MLFlow error details
- Workflow fails after retry attempts exhausted
Assertions:
- Exception raised from transform activity
- Error notification sent
- Workflow fails
- Export NOT called
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 207"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(207, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(207, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock request_transform to raise an exception
def mock_request_transform(*args, **kwargs):
raise Exception("MLFlow transform service unavailable")
with patch.object(test_activities, 'request_transform', side_effect=mock_request_transform):
input_data = {
'metadata': {
'metadata': {
'model_id': 207,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 207,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 207',
'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'],
}
print("\n[TEST] 2. Starting workflow that should fail on transform...")
workflow_id = f'test-transform-error-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow to fail...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
pytest.fail("Expected workflow to fail, but it completed successfully")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
# Verify no predictions were created
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 207")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_2_2_mlflow_predict_api_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 2.4.2: MLFlow Predict API Error
Description:
MLFlow predict request fails.
Expected Behavior:
- request_predict raises exception
- Notification sent
- Workflow fails after retries
Assertions:
- Transform succeeded
- Predict raised exception
- Error notification sent
- Workflow fails
"""
client = temporal_test_env.client
print("\n[TEST] 1. Inserting test data...")
with postgres_engine.begin() as conn:
conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 208"))
insert_sql = """
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
VALUES
(208, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
(208, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
"""
conn.execute(text(insert_sql))
print("[TEST] ✓ Data inserted successfully")
# Mock request_predict to raise an exception
def mock_request_predict(*args, **kwargs):
raise Exception("MLFlow predict service unavailable")
with patch.object(test_activities, 'request_predict', side_effect=mock_request_predict):
input_data = {
'metadata': {
'metadata': {
'model_id': 208,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}
},
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': 208,
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 208',
'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'],
}
print("\n[TEST] 2. Starting workflow that should fail on predict...")
workflow_id = f'test-predict-error-{datetime.now().timestamp()}'
handle = await client.start_workflow(
PredictionsBatch.run,
input_data,
id=workflow_id,
task_queue='test-queue',
)
print("[TEST] ✓ Workflow started")
print("\n[TEST] 3. Waiting for workflow to fail...")
try:
await asyncio.wait_for(handle.result(), timeout=60.0)
pytest.fail("Expected workflow to fail, but it completed successfully")
except Exception as e:
print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
# Verify no predictions were created
with postgres_engine.connect() as conn:
result_query = conn.execute(
text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 208")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
print("\n[TEST] ✓ All assertions passed!")