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

SIENTIAPDE-1478

Browse Source 
Enhance end-to-end tests for PredictionsBatch workflow scenarios

- Introduced mock repositories for PI Web API and OPC operations to improve test coverage.
- Updated test scenarios to handle partial write errors for PI Web API and OPC.
- Refactored existing tests to assert correct behavior under various error conditions.
- Enhanced logging and error handling in API and OPC activities to provide clearer feedback on failures.
- Removed outdated integration test file to streamline test suite.
This commit is contained in:
vitor-aignosi
2026-01-16 16:41:36 -03:00
parent 6b9bb38968
commit 241f283724
9 changed files with 1546 additions and 547 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
e2e/conftest.py
View File

@@ -225,6 +225,32 @@ def mock_minio_repository():
return mock_repo
@pytest_asyncio.fixture
def mock_pi_web_api_repository():
"""Mock PI Web API repository for PI Web API operations."""
mock_repo = MagicMock()
mock_repo.write_value = AsyncMock(
return_value={
'Items': [
{
'WebId': 'web_id_1'
}
]
}
)
mock_repo.close = MagicMock()
return mock_repo
@pytest_asyncio.fixture
def mock_opc_repository():
"""Mock OPC repository for OPC operations."""
mock_repo = MagicMock()
mock_repo.write_data = AsyncMock(
return_value=(True, {'response_time': 0.1})
)
mock_repo.disconnect = MagicMock()
return mock_repo
@pytest_asyncio.fixture
def patch_create_engine(postgres_engine):
"""Patch create_engine to return test postgres_engine."""
@@ -238,6 +264,11 @@ def patch_minio_repository(mock_minio_repository):
with patch('laborious.utils.repository.minio_repository.MinioRepository', return_value=mock_minio_repository):
yield
@pytest_asyncio.fixture
def patch_pi_web_api_repository(mock_pi_web_api_repository):
"""Patch MLflowRepository to return mock."""
with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
yield
@pytest_asyncio.fixture
def mock_mlflow_models():
@@ -330,6 +361,8 @@ async def test_activities(
patch_create_engine,
patch_minio_repository,
patch_mlflow,
patch_pi_web_api_repository,
mock_opc_repository
):
"""
Create Activities instance with test dependencies.
@@ -372,6 +405,10 @@ async def test_activities(
notification_handler=notification_handler,
)
activities.opc_repository = {
'1': mock_opc_repository,
}
try:
yield activities
finally:

92
e2e/scenarios.md
View File

@@ -445,27 +445,7 @@ The `predictions_batch` workflow:
### 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
#### Scenario 3.2.1: PI Web API Write Error
**Description**: PI Web API export fails
**Input**:
@@ -485,7 +465,7 @@ The `predictions_batch` workflow:
---
#### Scenario 3.2.3: OPC Write Error
#### Scenario 3.2.2: OPC Write Error
**Description**: OPC server write fails
**Input**:
@@ -504,64 +484,26 @@ The `predictions_batch` workflow:
---
## 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
#### Scenario 3.2.3: PI Web API Partial Write Error
**Description**: Two prediction tags attempt to be written to PI Web API, but only one succeeds
**Input**:
- Valid SQL query returning data
- All configurations provided (OPC, PI Web API, filters, policies)
- MLFlow services available
- All databases available
- Valid prediction
- Two prediction tags configured
- PI Web API returns partial success (one tag succeeds, one fails)
**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
- `write_pi_web_api_data` processes response
- `process_pi_web_api_response` detects partial failure
- Error confidence set (13)
- Notification sent for failed tag
- Workflow completes with error confidence
**Assertions**:
- 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
- One tag written successfully
- One tag failed
- Error confidence set in prediction
- Error notification sent
- Workflow completes
---

586
e2e/test_predictions_batch_format_export.py
View File

@@ -4,10 +4,11 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
import asyncio
from datetime import datetime
from unittest.mock import patch
from unittest.mock import ANY, AsyncMock, patch, call
import pandas as pd
import pytest
from sientia_do.notifications.models import NotificationLevel
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
@@ -87,6 +88,37 @@ async def start_and_await_workflow(client, input_data, workflow_id):
except asyncio.TimeoutError:
pytest.fail("Workflow execution timed out after 60 seconds")
def assert_prediction(
postgres_engine, model_id, prediction: float = 0.5,
prediction_confidence: int = 0, prediction_status: str = 'Good',
comments: str = '',
):
"""
Verify prediction was created with correct values in database
Args:
postgres_engine: Database engine
model_id: Model ID to check
prediction: Expected prediction value (default 0.5 from mock)
prediction_confidence: Expected confidence value (default 0 for normal predictions)
prediction_status: Expected status (default 'Good')
comments: Expected comments (default empty string)
"""
print("\n[TEST] 4. Verifying prediction was created with correct values...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, f"Expected one prediction record, got {len(prediction_rows)}"
row = prediction_rows[0]
assert row[0] == model_id, f"Expected model_id={model_id}, got {row[0]}"
assert row[1] == prediction, f"Expected prediction={prediction}, got {row[1]}"
assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}"
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
@pytest.mark.asyncio
@pytest.mark.integration
@@ -151,6 +183,63 @@ async def test_scenario_3_1_1_default_prediction_export(
await start_and_await_workflow(client, input_data, workflow_id)
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
call(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0.5,
},
endpoint='test_endpoint',
metadata={
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
},
endpoint='test_endpoint',
metadata={
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
],
any_order=True,
)
test_activities.opc_repository['1'].write_data.assert_has_calls(
[
call('addr_1', 0.5, 'float', ANY,
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
call('addr_2', 0, 'float', ANY,
{
'model_id': 311,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
]
)
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@@ -191,8 +280,18 @@ async def test_scenario_3_1_2_export_with_opc_only(
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = {
'server_name': 'test_server',
'tags': {'prediction': 'test_tag'},
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
input_data['pi_web_api_output_config'] = None # No PI Web API config
@@ -201,14 +300,29 @@ async def test_scenario_3_1_2_export_with_opc_only(
await start_and_await_workflow(client, input_data, workflow_id)
print("\n[TEST] 4. Verifying PostgreSQL and OPC export were executed...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
test_activities.opc_repository['1'].write_data.assert_has_calls(
[
call('addr_1', 0.5, 'float', ANY,
{
'model_id': 312,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
call('addr_2', 0, 'float', ANY,
{
'model_id': 312,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
]
)
test_activities.pi_web_api_client.write_value.assert_not_called()
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@@ -250,8 +364,8 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
input_data = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'prediction': 'test_pred_tag'},
'confidence_tags': {'confidence': 'test_conf_tag'},
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = None # No OPC config
@@ -260,14 +374,44 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
await start_and_await_workflow(client, input_data, workflow_id)
print("\n[TEST] 4. Verifying PostgreSQL and PI Web API export were executed...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
call(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0.5,
},
endpoint='test_endpoint',
metadata={
'model_id': 313,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
},
endpoint='test_endpoint',
metadata={
'model_id': 313,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
],
any_order=True,
)
test_activities.opc_repository['1'].write_data.assert_not_called()
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@@ -314,14 +458,11 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
await start_and_await_workflow(client, input_data, workflow_id)
print("\n[TEST] 4. Verifying only PostgreSQL export was executed...")
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, "Expected one prediction record in PostgreSQL"
test_activities.pi_web_api_client.write_value.assert_not_called()
test_activities.opc_repository['1'].write_data.assert_not_called()
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@@ -361,116 +502,106 @@ async def test_scenario_3_1_5_export_without_transformed_data(
input_data = get_base_input_data(model_id)
input_data['save_transform'] = False # Don't save transformed data
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
print("\n[TEST] 2. Starting workflow without transformed data export...")
workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
print("\n[TEST] 4. Verifying only prediction was exported...")
test_activities.pi_web_api_client.write_value.assert_has_calls(
[
call(
web_ids=['web_id_1'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0.5,
},
endpoint='test_endpoint',
metadata={
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
call(
web_ids=['web_id_2'],
value={
'Timestamp': '2024-01-01 12:00:00+0000',
'Value': 0,
},
endpoint='test_endpoint',
metadata={
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
},
),
],
any_order=True,
)
test_activities.opc_repository['1'].write_data.assert_has_calls(
[
call('addr_1', 0.5, 'float', ANY,
{
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
call('addr_2', 0, 'float', ANY,
{
'model_id': 315,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'predictions_batch',
}),
]
)
with postgres_engine.connect() as conn:
# Verify prediction exists
result_query = conn.execute(
text(f"SELECT model_id FROM predictions_schema.predictions WHERE model_id = {model_id}")
)
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(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
)
count = result_query.scalar()
assert count == 0, f"Expected transform table to be empty, but found {count} records"
assert_prediction(postgres_engine, model_id)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_1_postgres_export_error_predictions_table(
async def test_scenario_3_2_1_pi_web_api_write_error(
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
model_id = 321
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
# Mock export_data_to_postgres to raise an exception
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 = get_base_input_data(model_id)
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(f"SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}")
)
count = result_query.scalar()
assert count == 0, f"Expected no predictions, but found {count} records"
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
Scenario 3.2.1: PI Web API Write Error
Description:
PI Web API export fails.
@@ -488,55 +619,60 @@ async def test_scenario_3_2_2_pi_web_api_write_error(
"""
client = temporal_test_env.client
model_id = 322
model_id = 321
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
# Mock write_pi_web_api_data to raise an exception
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 = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'prediction': 'test_pred_tag'},
'confidence_tags': {'confidence': 'test_conf_tag'},
test_activities.pi_web_api_client.write_value.side_effect = Exception(
"PI Web API service unavailable")
input_data = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
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] 2. Starting workflow that should fail on PI Web API write...")
workflow_id = f'test-pi-api-error-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
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!")
assert_prediction(
postgres_engine, model_id,
prediction_confidence=13,
comments='PI Web API service unavailable',
)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_3_opc_write_error(
async def test_scenario_3_2_2_opc_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.3: OPC Write Error
Scenario 3.2.2: OPC Write Error
Description:
OPC server write fails.
@@ -553,39 +689,141 @@ async def test_scenario_3_2_3_opc_write_error(
"""
client = temporal_test_env.client
model_id = 322
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
test_activities.opc_repository['1'].write_data.return_value = (False, {
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
'message': 'OPC server unavailable',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'OPC server unavailable',
})
input_data = get_base_input_data(model_id)
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
print("\n[TEST] 2. Starting workflow that should fail on OPC write...")
workflow_id = f'test-opc-error-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
assert_prediction(
postgres_engine, model_id,
prediction_confidence=12,
comments='Some data could not be written to OPC servers',
)
print("\n[TEST] ✓ All assertions passed!")
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_3_2_3_pi_web_api_partial_write_error(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
):
"""
Scenario 3.2.3: PI Web API Partial Write Error
Description:
Two prediction tags attempt to be written to PI Web API, but only one succeeds.
Expected Behavior:
- write_pi_web_api_data processes response
- process_pi_web_api_response detects partial failure
- Error confidence set (13)
- Notification sent for failed tag
- Workflow completes with error confidence
Assertions:
- One tag written successfully
- One tag failed
- Error confidence set in prediction
- Error notification sent
- Workflow completes
"""
client = temporal_test_env.client
model_id = 323
print("\n[TEST] 1. Inserting test data...")
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
print("[TEST] ✓ Data inserted successfully")
# Mock write_opc_data to raise an exception
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 = get_base_input_data(model_id)
input_data['opc_output_config'] = {
'server_name': 'test_server',
'tags': {'prediction': 'test_tag'},
test_activities.pi_web_api_client.write_value = AsyncMock(side_effect=[
{
'Items': [
{
'WebId': 'web_id_1',
'Errors': [],
},
]
},
Exception('Tag write failed'),
{
'Items': [
{
'WebId': 'web_id_2',
'Errors': [],
},
]
},
])
input_data = get_base_input_data(model_id)
input_data['pi_web_api_output_config'] = {
'endpoint': 'test_endpoint',
'prediction_tags': {'tag_1': 'web_id_1', 'tag_3': 'web_id_3'},
'confidence_tags': {'tag_2': 'web_id_2'},
}
input_data['opc_output_config'] = {
'1': {
'prediction_tags': {
'addr_1': {
'data_type': 'float',
}
},
'confidence_tags': {
'addr_2': {
'data_type': 'float',
}
},
}
}
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] 2. Starting workflow with partial PI Web API write error...")
workflow_id = f'test-pi-api-partial-error-{datetime.now().timestamp()}'
await start_and_await_workflow(client, input_data, workflow_id)
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!")
assert_prediction(
postgres_engine, model_id,
prediction_confidence=13,
comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.",
)
print("\n[TEST] ✓ All assertions passed!")

278
e2e/test_predictions_batch_integration.py
View File

@@ -1,278 +0,0 @@
"""
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!")

20
laborious/activities/api.py
View File

@@ -81,7 +81,7 @@ class API(SientiaMonitoring):
tags: dict[str, str],
core_labels: dict[str, str],
metadata: dict[str, Any],
) -> int:
) -> tuple[int, str]:
"""
Process the response data from PI Web API write operation.
@@ -106,6 +106,8 @@ class API(SientiaMonitoring):
confidence = 0
message = ''
# Evaluate response for each tag
written_tags = []
response_items = response_data.get('Items', [])
@@ -144,10 +146,10 @@ class API(SientiaMonitoring):
written_tags.append(tag_name)
if len(written_tags) != len(tag_names):
self.error(
f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written',
metadata,
)
message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.'
self.error(f"{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}", metadata)
await self.send_notification_async(
metadata=metadata,
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
@@ -157,7 +159,7 @@ class API(SientiaMonitoring):
)
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
return confidence
return confidence, message
@activity.defn(name='write_pi_web_api_data')
async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
@@ -185,7 +187,7 @@ class API(SientiaMonitoring):
data = DataFrame(input_data['data'])
pi_web_api_output_config = input_data['pi_web_api_output_config']
self.info('Writing data to PI Web API...', metadata)
self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata)
endpoint = pi_web_api_output_config['endpoint']
@@ -213,7 +215,7 @@ class API(SientiaMonitoring):
metadata=metadata,
)
confidence = await self.process_pi_web_api_response(
confidence, message = await self.process_pi_web_api_response(
response_data=prediction_response,
tags=raw_prediction_tags,
core_labels=core_labels,
@@ -221,6 +223,7 @@ class API(SientiaMonitoring):
)
data['prediction_confidence'] = confidence
data['comments'] = message
except Exception as e:
trace = traceback.format_exc()
@@ -234,6 +237,7 @@ class API(SientiaMonitoring):
)
data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
data['comments'] = str(e)
return data.to_dict()

5
laborious/activities/opc.py
View File

@@ -352,10 +352,13 @@ class OPC(SientiaMonitoring):
This allows downstream systems to handle data quality appropriately.
"""
message = 'Some data could not be written to OPC servers'
if not success:
data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
data['comments'] = message
self.debug(
f'Some data could not be written to OPC servers, setting confidence to {OPC_WRITTING_ERROR_CONFIDENCE}.',
f'{message}, setting confidence to {OPC_WRITTING_ERROR_CONFIDENCE}.',
metadata,
)

1050
tests.ipynb
View File

File diff suppressed because one or more lines are too long

21
tests/laborious/activities/test_api.py
View File

@@ -127,8 +127,6 @@ async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_dat
result = await api.write_pi_web_api_data(input_data)
api.info.assert_called_once_with('Writing data to PI Web API...', metadata['metadata'])
api.pi_web_api_client.write_value.assert_has_calls(
[
call(
@@ -292,7 +290,7 @@ async def test_process_pi_web_api_response_success(api):
'workflow_name': 'test_workflow',
}
confidence = await api.process_pi_web_api_response(
confidence, message = await api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
@@ -300,6 +298,7 @@ async def test_process_pi_web_api_response_success(api):
)
assert confidence == 0
assert message == ''
assert api.emit_metric.call_count == 2
# Verify that emit_metric was called with correct tags structure
call_args_list = api.emit_metric.call_args_list
@@ -326,7 +325,7 @@ async def test_process_pi_web_api_response_with_errors(api):
'workflow_name': 'test_workflow',
}
confidence = await api.process_pi_web_api_response(
confidence, message = await api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
@@ -334,10 +333,9 @@ async def test_process_pi_web_api_response_with_errors(api):
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert message == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
assert api.emit_metric.call_count == 2
api.error.assert_any_call(
"Error writing tag tag1:web_id_1 to PI Web API: ['Error writing tag']", metadata['metadata']
)
@mark.asyncio
@@ -355,7 +353,7 @@ async def test_process_pi_web_api_response_missing_tags(api):
'workflow_name': 'test_workflow',
}
confidence = await api.process_pi_web_api_response(
confidence, message = await api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
@@ -363,6 +361,7 @@ async def test_process_pi_web_api_response_missing_tags(api):
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert message == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
api.send_notification_async.assert_called_once()
call_args = api.send_notification_async.call_args
assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
@@ -385,7 +384,7 @@ async def test_process_pi_web_api_response_missing_webid(api):
'workflow_name': 'test_workflow',
}
confidence = await api.process_pi_web_api_response(
confidence, message = await api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
@@ -393,6 +392,7 @@ async def test_process_pi_web_api_response_missing_webid(api):
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert message == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])
@@ -411,7 +411,7 @@ async def test_process_pi_web_api_response_missing_tag_name(api):
'workflow_name': 'test_workflow',
}
confidence = await api.process_pi_web_api_response(
confidence, message = await api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
@@ -419,6 +419,7 @@ async def test_process_pi_web_api_response_missing_tag_name(api):
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert message == "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
api.error.assert_any_call(
'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
)

4
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -812,7 +812,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
data = MagicMock()
output = await mlflow_repository.fit_models(
'model_name', data, 'latest_production_id', metadata['metadata'], 'sklearn', 'pyfunc', None
'model_name', data, 'latest_production_id', metadata['metadata'], 'sklearn', False, 'pyfunc', None
)
mlflow_repository.download_model.assert_has_calls(
@@ -908,6 +908,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
'latest_production_id',
metadata['metadata'],
'sklearn',
False,
'pyfunc',
'feat_1',
)
@@ -1425,6 +1426,7 @@ async def test_retrain_model(mlflow_repository):
model_name=model_name,
data=data,
transform_flavor='sklearn',
skip_transform=False,
predict_flavor='pyfunc',
target_name='target',
metadata=metadata['metadata'],