sientia-dataops-laborious-legacy/e2e/test_predictions_batch_prediction_process.py

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