sientia-dataops-laborious_temporal/e2e/test_predictions_batch_integration.py

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