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