""" 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!")