vitor-aignosi
503d9aa485
Update dependencies and refactor input filter handling for consistency - Updated sientia-dataops-library dependency version from 1.10.3 to 1.10.4 in requirements.txt. - Refactored input filter handling in the Gates class to read policy and config keys in a case-insensitive manner. - Updated test cases to ensure consistency in filter key naming conventions across various scenarios.
451 lines
17 KiB
Python
451 lines
17 KiB
Python
"""
|
|
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!")
|