sientia-dataops-laborious_temporal/e2e/conftest.py

"""
Pytest configuration and fixtures for E2E tests.
"""

import sys
from unittest.mock import AsyncMock, MagicMock, patch

# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
_model_analysis_module = MagicMock()
_model_analysis_module.ModelAnalysis = MagicMock
sys.modules.setdefault('sientia', MagicMock())
sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
from io import BytesIO

import joblib
import pandas as pd
import pytest
import pytest_asyncio
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.metrics_controller import MetricsController
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sqlalchemy import create_engine, text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from testcontainers.minio import MinioContainer
from testcontainers.postgres import PostgresContainer

from e2e.opc_test_server import OpcE2ETestServer
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess

# Test constants
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'

# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised.
JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model'
JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact'


@pytest_asyncio.fixture(scope='session')
def minio_container():
    """
    MinIO S3-compatible storage for E2E tests that exercise real offload uploads.
    """
    minio = MinioContainer()
    minio.start()
    yield minio
    minio.stop()


@pytest_asyncio.fixture(scope='session')
def postgres_container():
    """
    Create a PostgreSQL container using testcontainers.

    This fixture creates a real PostgreSQL database in a Docker container
    that will be used for all tests in the session.
    """
    postgres = PostgresContainer('postgres:15')
    postgres.start()
    yield postgres
    postgres.stop()


@pytest_asyncio.fixture
def postgres_engine(postgres_container):
    """
    Create SQLAlchemy engine for PostgreSQL test database.

    This fixture creates a connection to the PostgreSQL container
    created by the postgres_container fixture.
    """
    engine = create_engine(postgres_container.get_connection_url())

    yield engine

    engine.dispose()


def _create_schema_and_tables(engine):
    """
    Helper function to create schema and tables in the given engine.

    Creates predictions_schema with:
    - laborious_data: Input data table for queries
    - predictions: Output predictions table
    - transformed_data: Output transformed data table
    """
    # Use begin() to ensure transaction is properly committed
    with engine.begin() as conn:
        # Create predictions_schema
        conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema'))

        # Create laborious_data table (input data from sensors)
        create_laborious_data_sql = """
            CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
                id SERIAL NOT NULL,
                model_id int4 NOT NULL,
                variable text NOT NULL,
                value numeric NULL,
                "timestamp" timestamptz NOT NULL,
                created_at timestamptz NOT NULL,
                PRIMARY KEY (id)
            );
        """
        conn.execute(text(create_laborious_data_sql))

        # Create predictions table
        create_predictions_sql = """
            CREATE TABLE if not exists predictions_schema.predictions (
                id SERIAL NOT NULL ,
                model_id int4 NOT NULL,
                prediction numeric NULL,
                prediction_confidence numeric NOT NULL,
                response_time numeric NOT NULL,
                prediction_status text NOT NULL,
                "timestamp" timestamptz NOT NULL,
                created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
                "comments" text NULL,
                PRIMARY KEY (id, created_at)
            );
        """
        conn.execute(text(create_predictions_sql))

        # Create transformed_data table
        create_transformed_sql = """
            CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
                id SERIAL NOT NULL,
                model_id int4 NOT NULL,
                variable text NOT NULL,
                value numeric NULL,
                "timestamp" timestamptz NOT NULL,
                created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
                PRIMARY KEY (id)
            );
        """
        conn.execute(text(create_transformed_sql))


@pytest_asyncio.fixture(autouse=True)
def setup_postgres_schema_and_tables(postgres_engine):
    """
    Automatically create necessary schema and tables before each test.

    This fixture runs automatically (autouse=True) and ensures
    that the predictions_schema and tables exist with the correct structure.
    """
    _create_schema_and_tables(postgres_engine)
    yield


@pytest_asyncio.fixture
def mock_logger():
    """Mock logger for testing."""

    def message(message):
        print(f'[LOG] {message}')

    def custom_message(message, _metadata=None):
        print(f'[LOG] {message}')

    logger = MagicMock()
    logger.info = MagicMock(side_effect=message)
    logger.debug = MagicMock(side_effect=message)
    logger.error = MagicMock(side_effect=message)
    logger.warning = MagicMock(side_effect=message)
    logger.custom_info = MagicMock(side_effect=custom_message)
    logger.custom_debug = MagicMock(side_effect=custom_message)
    logger.custom_error = MagicMock(side_effect=custom_message)
    logger.custom_warning = MagicMock(side_effect=custom_message)
    return logger


@pytest_asyncio.fixture
def mock_mongo_client():
    """
    Mock MongoDB client to avoid real connections.

    This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
    allowing us to use a real NotificationHandler instance without connecting to MongoDB.
    """
    mock_client = MagicMock()
    mock_db = MagicMock()
    mock_collection = MagicMock()

    # Configure the mock chain: client[database] -> db[collection] -> collection
    mock_client.__getitem__.return_value = mock_db
    mock_db.__getitem__.return_value = mock_collection

    # Mock server_info() to avoid connection attempts
    mock_client.server_info = MagicMock()

    # Mock insert_one for notifications
    mock_collection.insert_one = MagicMock()

    return mock_client


@pytest.fixture
def notification_inserts(mock_mongo_client):
    """
    Mongo insert_one mock used by CoreNotificationHandler for notification persistence.

    Yields:
        MagicMock for insert_one, reset before each test.
    """
    mock_db = mock_mongo_client.__getitem__.return_value
    mock_collection = mock_db.__getitem__.return_value
    mock_collection.insert_one.reset_mock()
    yield mock_collection.insert_one


@pytest_asyncio.fixture
def notification_handler(mock_logger, mock_mongo_client):
    """
    Create a real NotificationHandler instance with mocked MongoDB client.

    This fixture creates a real CoreNotificationHandler instance but mocks
    the underlying MongoDB connection to avoid real database connections.
    """
    # Patch MongoClient where it's imported in the handlers module
    with patch('sientia_do.notifications.handlers.MongoClient', return_value=mock_mongo_client):
        handler = CoreNotificationHandler(
            connection_string=TEST_MONGODB_CONNECTION_STRING,
            database=TEST_DATABASE_NAME,
            logger=mock_logger,
            project_name='laborious',
        )
        yield handler
        handler.shutdown()


@pytest_asyncio.fixture
def metrics_controller(mock_logger):
    """Create a real MetricsController instance."""
    return MetricsController(logger=mock_logger)


@pytest_asyncio.fixture
def mock_minio_repository():
    """Mock MinIO repository for object storage operations."""
    mock_repo = MagicMock()

    # Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve()
    # can decode the payload if offloading is exercised in an integration scenario.
    parquet_df = pd.DataFrame({'a': [1]})
    parquet_buffer = BytesIO()
    parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
    parquet_bytes = parquet_buffer.getvalue()

    # sientia_do MinioRepository API
    mock_repo.bucket = 'test-bucket'
    mock_repo.upload_file = AsyncMock(
        side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: {
            'minio_object_name': f'sientia/streamlit-connectors/{relative_key}',
            'original_filename': relative_key.rsplit('/', 1)[-1],
            'uploaded_at': '2024-01-01T00:00:00Z',
            'sha256_hash': 'deadbeef',
        }
    )
    mock_repo.download_file = AsyncMock(return_value=parquet_bytes)
    mock_repo.list_objects = AsyncMock(return_value=[])
    mock_repo.delete_file = AsyncMock()
    mock_repo.close = MagicMock()

    return mock_repo


@pytest_asyncio.fixture
def mock_pi_web_api_repository():
    """Mock PI Web API repository for PI Web API operations."""
    mock_repo = MagicMock()

    async def _write_value(web_ids, value, metadata=None, **kwargs):
        """
        Mirror successful PI writes: one response item per requested web_id.

        write_pi_web_api_data passes the list into process_pi_web_api_response (not a
        wrapped {'Items': ...} envelope).
        """
        return [{'WebId': wid, 'Errors': []} for wid in web_ids]

    mock_repo.write_value = AsyncMock(side_effect=_write_value)
    mock_repo.close = MagicMock()
    return mock_repo


@pytest_asyncio.fixture
def mock_opc_repository():
    """Mock OPC repository for OPC operations."""
    mock_repo = MagicMock()
    mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1}))
    mock_repo.disconnect = AsyncMock()
    return mock_repo


@pytest_asyncio.fixture
async def opc_e2e_server():
    """
    In-process asyncua OPC UA server for E2E tests against OpcRepository.
    """
    server = OpcE2ETestServer()
    await server.start()
    try:
        yield server
    finally:
        await server.stop()


@pytest_asyncio.fixture
def patch_create_engine(postgres_engine):
    """Patch create_engine to return test postgres_engine."""
    with patch(
        'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine
    ):
        yield


@pytest_asyncio.fixture
def patch_minio_repository(mock_minio_repository):
    """Patch MinioRepository to return mock."""
    # Patch where Activities resolves the symbol (import binds the original class).
    with patch(
        'laborious.activities.activities.MinioRepository', return_value=mock_minio_repository
    ):
        yield


@pytest_asyncio.fixture
def patch_pi_web_api_repository(mock_pi_web_api_repository):
    """Patch MLflowRepository to return mock."""
    with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
        yield


@pytest_asyncio.fixture
def mock_mlflow_models():
    """Create mock models for MLflow load_model methods."""
    # Mock transform model - returns DataFrame with same index as input
    mock_transform_model = MagicMock()

    def mock_transform_predict(data):
        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
        print(data.to_csv())
        print(data.index)
        result = pd.DataFrame(
            {
                'feature_1': [0.234] * num_rows,
                'feature_2': [0.783] * num_rows,
            }
        )
        result.index = data.index
        return result

    mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)

    # Mock predict model - returns array/list of predictions
    mock_predict_model = MagicMock()

    def mock_predict_predict(data):
        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
        return [0.5] * num_rows

    mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)

    # Mock PyFuncModel for compressed models
    mock_pyfunc_model = MagicMock()
    mock_pyfunc_model._model_impl = MagicMock()
    mock_pyfunc_model._model_impl.python_model = mock_transform_model

    return {
        'transform_model': mock_transform_model,
        'predict_model': mock_predict_model,
        'pyfunc_model': mock_pyfunc_model,
    }


def _build_joblib_fallback_estimator():
    """
    Build a real sklearn estimator for the joblib fallback artifact.

    `request_predict` appends a string `timestamp` column to the transformed
    data, so the estimator selects the transform output columns explicitly and
    drops the rest.
    """
    features = ['feature_1', 'feature_2']
    estimator = Pipeline(
        steps=[
            (
                'select',
                ColumnTransformer([('features', 'passthrough', features)], remainder='drop'),
            ),
            ('regressor', LinearRegression()),
        ]
    )
    train = pd.DataFrame(
        {
            'feature_1': [0.0, 1.0, 2.0, 3.0],
            'feature_2': [0.0, 1.0, 0.0, 1.0],
            'timestamp': ['2024-01-01 12:00:00'] * 4,
        }
    )
    estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3)
    return estimator


@pytest.fixture
def joblib_fallback_artifact(tmp_path):
    """
    Real artifact directory holding a joblib-compressed sklearn estimator.

    Mirrors what the joblib fallback finds after downloading `prediction_model`:
    a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an
    `MLmodel` placeholder.

    Return:
        dict with 'path' (the artifact directory) and 'estimator' (the object
        that was serialized, so tests can compute the expected prediction).
    """
    artifact_dir = tmp_path / 'prediction_model'
    artifact_dir.mkdir()

    estimator = _build_joblib_fallback_estimator()
    joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3)
    (artifact_dir / 'MLmodel').write_text('flavors:\n  sklearn: {}\n')

    return {'path': str(artifact_dir), 'estimator': estimator}


@pytest.fixture
def joblib_fallback_artifact_without_pickle(tmp_path):
    """Artifact directory with no top-level `.pkl`, so the joblib fallback also fails."""
    artifact_dir = tmp_path / 'prediction_model_without_pickle'
    artifact_dir.mkdir()
    (artifact_dir / 'MLmodel').write_text('flavors:\n  sklearn: {}\n')
    return str(artifact_dir)


@pytest_asyncio.fixture
def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
    """Patch mlflow module in repository with load_model mocks."""
    mock_mlflow = MagicMock()

    # Mock sklearn.load_model
    def mock_sklearn_load_model(model_uri):
        if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
            raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
        if 'data_model' in model_uri or 'transform' in model_uri.lower():
            return mock_mlflow_models['transform_model']
        return mock_mlflow_models['predict_model']

    mock_mlflow.sklearn = MagicMock()
    mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)

    # Mock pyfunc.load_model
    def mock_pyfunc_load_model(model_uri):
        if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
            raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
        if 'artifacts' in model_uri or 'tmp' in model_uri:
            return mock_mlflow_models['pyfunc_model']
        if 'data_model' in model_uri or 'transform' in model_uri.lower():
            return mock_mlflow_models['transform_model']
        return mock_mlflow_models['predict_model']

    mock_mlflow.pyfunc = MagicMock()
    mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)

    # Mock pytorch.load_model
    mock_mlflow.pytorch = MagicMock()
    mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])

    # Mock other mlflow methods that might be called
    mock_mlflow.set_tracking_uri = MagicMock()
    mock_mlflow.get_run = MagicMock(
        return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))
    )
    mock_mlflow.tracking = MagicMock()
    mock_mlflow_client = MagicMock(
        search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
        search_model_versions=MagicMock(
            return_value=[
                MagicMock(
                    current_stage='Production', version='1', source='runs:/artifacts/test_run_id'
                )
            ]
        ),
    )
    # Only the joblib fallback calls download_artifacts directly on the client.
    mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path'])
    mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)

    with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
        yield mock_mlflow


@pytest.fixture
def mlflow_client(patch_mlflow):
    """The MlflowClient mock the repository holds as `self.client`."""
    return patch_mlflow.tracking.MlflowClient.return_value


@pytest_asyncio.fixture(scope='function')
async def test_activities(
    postgres_engine,
    postgres_container,
    mock_logger,
    notification_handler,
    metrics_controller,
    mock_minio_repository,
    patch_create_engine,
    patch_minio_repository,
    patch_mlflow,
    patch_pi_web_api_repository,
    mock_opc_repository,
):
    """
    Create Activities instance with test dependencies.

    This fixture creates a real Activities instance with:
    - PostgreSQL database (via testcontainers)
    - Mocked MinIO client
    - Real NotificationHandler and MetricsController (with mocked underlying services)
    """
    activities = Activities(
        postgres_config={
            'host': 'localhost',
            'port': postgres_container.get_exposed_port(5432),
            'user': 'test',
            'password': 'test',
            'dbname': 'test',
            'min_connections': 1,
            'max_connections': 5,
        },
        mlflow_config={
            'host': 'http://localhost',
            'port': '5000',
            'username': 'test',
            'password': 'test',
        },
        minio_config={
            # Host:port only; Minio() prepends http(s):// from the secure flag.
            'endpoint_url': 'localhost:9000',
            'access_key': 'test',
            'secret_key': 'test',
            'default_bucket': 'test-bucket',
            'retention_hours': 24,
            'secure': False,
        },
        opc_config={},
        pi_web_api_config={
            'base_url': 'http://localhost:8080',
            'auth_type': 'bearer',
            'auth_token': 'test_token',
        },
        logger=mock_logger,
        notification_handler=notification_handler,
    )

    activities.opc_repository = {
        '1': mock_opc_repository,
    }

    try:
        yield activities
    finally:
        # Cleanup - ALWAYS runs, even if test fails
        await activities.shutdown()


@pytest_asyncio.fixture(scope='function')
async def test_activities_real_minio(
    postgres_engine,
    postgres_container,
    minio_container,
    mock_logger,
    notification_handler,
    metrics_controller,
    patch_create_engine,
    patch_mlflow,
    patch_pi_web_api_repository,
    mock_opc_repository,
):
    """
    Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests.
    """
    minio_client = minio_container.get_client()
    if not minio_client.bucket_exists('test-bucket'):
        minio_client.make_bucket('test-bucket')
    minio_port = minio_container.get_exposed_port(9000)
    activities = Activities(
        postgres_config={
            'host': 'localhost',
            'port': postgres_container.get_exposed_port(5432),
            'user': 'test',
            'password': 'test',
            'dbname': 'test',
            'min_connections': 1,
            'max_connections': 5,
        },
        mlflow_config={
            'host': 'http://localhost',
            'port': '5000',
            'username': 'test',
            'password': 'test',
        },
        minio_config={
            'endpoint_url': f'localhost:{minio_port}',
            'access_key': 'minioadmin',
            'secret_key': 'minioadmin',
            'default_bucket': 'test-bucket',
            'retention_hours': 24,
            'secure': False,
        },
        opc_config={},
        pi_web_api_config={
            'base_url': 'http://localhost:8080',
            'auth_type': 'bearer',
            'auth_token': 'test_token',
        },
        logger=mock_logger,
        notification_handler=notification_handler,
    )
    activities.opc_repository = {'1': mock_opc_repository}
    try:
        yield activities
    finally:
        await activities.shutdown()


def _worker_activity_list(test_activities: Activities):
    return [
        test_activities.load_custom_query,
        test_activities.load_query_with_minio_offload,
        test_activities.cleanup_minio_objects_expired,
        test_activities.input_gate,
        test_activities.request_transform,
        test_activities.mlflow_response_gate,
        test_activities.mlflow_content_gate,
        test_activities.request_predict,
        test_activities.repeat_last_prediction,
        test_activities.format_prediction,
        test_activities.format_transformed_data,
        test_activities.format_default_prediction,
        test_activities.write_pi_web_api_data,
        test_activities.write_opc_data,
        test_activities.export_data_to_postgres,
        test_activities.export_payload_to_postgres,
        test_activities.write_metrics,
    ]


@pytest_asyncio.fixture(scope='function')
async def temporal_test_env():
    """Create Temporal test environment."""
    env = await WorkflowEnvironment.start_time_skipping()
    async with env:
        yield env


@pytest_asyncio.fixture(scope='function')
async def temporal_worker(temporal_test_env, test_activities):
    """Create Temporal worker with test activities."""
    async with Worker(
        temporal_test_env.client,
        task_queue='test-queue',
        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
        activities=_worker_activity_list(test_activities),
    ) as worker:
        yield worker


@pytest_asyncio.fixture(scope='function')
async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio):
    """Temporal worker backed by Activities using real MinIO testcontainer."""
    async with Worker(
        temporal_test_env.client,
        task_queue='test-queue',
        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
        activities=_worker_activity_list(test_activities_real_minio),
    ) as worker:
        yield worker


@pytest_asyncio.fixture(scope='function')
async def test_activities_real_opc(
    postgres_engine,
    postgres_container,
    opc_e2e_server: OpcE2ETestServer,
    mock_logger,
    notification_handler,
    metrics_controller,
    patch_create_engine,
    patch_minio_repository,
    patch_mlflow,
    patch_pi_web_api_repository,
):
    """
    Activities with a real OpcRepository connected to the in-process OPC UA server.
    """
    activities = Activities(
        postgres_config={
            'host': 'localhost',
            'port': postgres_container.get_exposed_port(5432),
            'user': 'test',
            'password': 'test',
            'dbname': 'test',
            'min_connections': 1,
            'max_connections': 5,
        },
        mlflow_config={
            'host': 'http://localhost',
            'port': '5000',
            'username': 'test',
            'password': 'test',
        },
        minio_config={
            'endpoint_url': 'localhost:9000',
            'access_key': 'test',
            'secret_key': 'test',
            'default_bucket': 'test-bucket',
            'retention_hours': 24,
            'secure': False,
        },
        opc_config={
            '1': {
                'id': '1',
                'server_name': 'e2e-opc',
                'url': opc_e2e_server.url,
                'server_uri': opc_e2e_server.url,
                'cert_path': None,
                'private_key_path': None,
                'server_cert_path': None,
                'reconnection_interval': 0,
            }
        },
        pi_web_api_config={
            'base_url': 'http://localhost:8080',
            'auth_type': 'bearer',
            'auth_token': 'test_token',
        },
        logger=mock_logger,
        notification_handler=notification_handler,
    )
    await activities.init_opc()
    repo = activities.opc_repository['1']
    assert repo._session_ready.is_set(), 'OPC E2E server connection failed during init_opc'
    try:
        yield activities
    finally:
        await activities.shutdown()


@pytest_asyncio.fixture(scope='function')
async def temporal_worker_real_opc(temporal_test_env, test_activities_real_opc):
    """Temporal worker backed by Activities using the in-process OPC UA server."""
    async with Worker(
        temporal_test_env.client,
        task_queue='test-queue',
        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
        activities=_worker_activity_list(test_activities_real_opc),
    ) as worker:
        yield worker