From 63017754a8e5c995837a199438b7e9e6f0b23d7e Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 14 Jan 2026 16:01:03 -0300 Subject: [PATCH] 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. --- e2e/__init__.py | 3 + e2e/conftest.py | 415 ++++++++++ e2e/scenarios.md | 561 ++++++++++++++ e2e/test_predictions_batch_format_export.py | 710 ++++++++++++++++++ e2e/test_predictions_batch_integration.py | 278 +++++++ e2e/test_predictions_batch_main_workflow.py | 450 +++++++++++ ...st_predictions_batch_prediction_process.py | 573 ++++++++++++++ laborious/activities/mlflow.py | 3 +- .../format_and_export_prediction.py | 4 +- pyproject.toml | 2 +- requirements-dev.txt | 3 +- tests/laborious/activities/test_gates.py | 1 + values.yaml | 6 +- 13 files changed, 3001 insertions(+), 8 deletions(-) create mode 100644 e2e/__init__.py create mode 100644 e2e/conftest.py create mode 100644 e2e/scenarios.md create mode 100644 e2e/test_predictions_batch_format_export.py create mode 100644 e2e/test_predictions_batch_integration.py create mode 100644 e2e/test_predictions_batch_main_workflow.py create mode 100644 e2e/test_predictions_batch_prediction_process.py diff --git a/e2e/__init__.py b/e2e/__init__.py new file mode 100644 index 0000000..b7a5d65 --- /dev/null +++ b/e2e/__init__.py @@ -0,0 +1,3 @@ +""" +End-to-end tests for laborious temporal workflows. +""" diff --git a/e2e/conftest.py b/e2e/conftest.py new file mode 100644 index 0000000..dd821ae --- /dev/null +++ b/e2e/conftest.py @@ -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 diff --git a/e2e/scenarios.md b/e2e/scenarios.md new file mode 100644 index 0000000..f9f8bbc --- /dev/null +++ b/e2e/scenarios.md @@ -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 + +--- \ No newline at end of file diff --git a/e2e/test_predictions_batch_format_export.py b/e2e/test_predictions_batch_format_export.py new file mode 100644 index 0000000..901d15a --- /dev/null +++ b/e2e/test_predictions_batch_format_export.py @@ -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!") diff --git a/e2e/test_predictions_batch_integration.py b/e2e/test_predictions_batch_integration.py new file mode 100644 index 0000000..28eab96 --- /dev/null +++ b/e2e/test_predictions_batch_integration.py @@ -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!") diff --git a/e2e/test_predictions_batch_main_workflow.py b/e2e/test_predictions_batch_main_workflow.py new file mode 100644 index 0000000..f8469ef --- /dev/null +++ b/e2e/test_predictions_batch_main_workflow.py @@ -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!") diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py new file mode 100644 index 0000000..03ff24b --- /dev/null +++ b/e2e/test_predictions_batch_prediction_process.py @@ -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!") diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 2e20cda..1847ff1 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -153,8 +153,9 @@ class MLFlow(SientiaMonitoring): # Pivot data for model input format data = data.pivot(index='timestamp', columns='variable', values='value') data.fillna(np.nan, inplace=True) - # data.reset_index(inplace=True) + data.columns.name = None + data.index.name = None self.debug('Processed input data:', metadata) self.debug(data.head(5).to_string(), metadata) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 6663f61..684f7a5 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -162,7 +162,7 @@ class FormatAndExportPrediction: opc_metrics = {} # write to pi web api - if pi_web_api_output_config is not None: + if pi_web_api_output_config: prediction = await workflow.execute_activity_method( Activities.write_pi_web_api_data, { @@ -175,7 +175,7 @@ class FormatAndExportPrediction: ) # write to opc - if opc_output_config is not None: + if opc_output_config: prediction, opc_metrics = await workflow.execute_activity_method( Activities.write_opc_data, { diff --git a/pyproject.toml b/pyproject.toml index 0f912c2..daaeabe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,7 +123,7 @@ markers = [ ] [tool.coverage.run] -source = ["model_manager"] +source = ["laborious"] omit = [ "*/tests/*", "*/venv/*", diff --git a/requirements-dev.txt b/requirements-dev.txt index 56ab376..16a8492 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -13,7 +13,8 @@ types-requests>=2.31.0 # Type stubs for requests pytest>=7.4.0 # Testing framework pytest-cov>=4.1.0 # Coverage plugin for pytest pytest-asyncio>=0.21.0 # Async test support (already in main requirements) +testcontainers[postgres] # PostgreSQL containers for E2E tests # Development Tools ipython>=8.12.0 # Enhanced Python shell -ipdb>=0.13.13 # IPython debugger \ No newline at end of file +ipdb>=0.13.13 # IPython debugger diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index e33ea4f..f124c30 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -595,6 +595,7 @@ async def test_format_transformed_data_multiple_rows(gates_activity): assert len(result['variable']) == 4 assert len(result['value']) == 4 assert len(result['model_id']) == 4 + assert len(result['created_at']) == 4 assert all(v == 'test_model' for v in result['model_id'].values()) assert set(result['variable'].values()) == {'var1', 'var2'} gates_activity.info.assert_called() diff --git a/values.yaml b/values.yaml index 87a56b6..5b19e15 100644 --- a/values.yaml +++ b/values.yaml @@ -11,7 +11,7 @@ image: # This sets the pull policy for images. pullPolicy: Always # Overrides the image tag whose default is the chart appVersion. - tag: "1.0.1" + tag: "1.1.0" # This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/ imagePullSecrets: @@ -71,7 +71,7 @@ livenessProbe: - -c - | curl -sf http://localhost:9090/metrics | grep -q '^app_up{.*} 1' - initialDelaySeconds: 360 + initialDelaySeconds: 660 periodSeconds: 15 timeoutSeconds: 5 failureThreshold: 3 @@ -83,7 +83,7 @@ readinessProbe: - -c - | curl -sf http://localhost:9090/metrics | grep -q '^app_up{.*} 1' - initialDelaySeconds: 300 + initialDelaySeconds: 600 periodSeconds: 10 timeoutSeconds: 3 failureThreshold: 2