"""Pytest configuration and fixtures for E2E tests.""" from concurrent.futures import ThreadPoolExecutor from typing import Any from unittest.mock import MagicMock import pandas as pd import pytest import pytest_asyncio from sqlalchemy import create_engine, text from testcontainers.core.container import DockerContainer from testcontainers.minio import MinioContainer from testcontainers.postgres import PostgresContainer from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker class _FakeModelAnalysis: """ Configurable double for ``sientia_model.analytics.model_analysis.ModelAnalysis`` used by E2E drift tests. Replaces the real analyzer in ``sientia_model.analytics.model_analysis`` BEFORE the production import chain runs so ``laborious.activities.model_metrics`` resolves ``ModelAnalysis`` to this class at import time. Keeps drift assertions stable across runs (the real analyzer is data-dependent). Tests configure responses through class-level attributes which are reset between tests by ``reset_model_analysis_stub``: - ``_drift_response_factory``: callable ``(univariate, multivariate) -> DataFrame`` controlling the consolidated drift dataframe seen by ``calculate_drift``. - ``_univariate_side_effect`` / ``_multivariate_side_effect``: optional side effects (Exception or callable) for the raw detection methods. - ``_drift_metrics_dataframe_exception``: when set, raised by ``get_drift_metrics_dataframe`` to simulate analyzer failures. """ _drift_response_factory: Any = None _univariate_side_effect: Any = None _multivariate_side_effect: Any = None _drift_metrics_dataframe_exception: Exception | None = None last_instance: Any = None def __init__(self, config): self.config = config type(self).last_instance = self self.detect_univariate_drift_calls = [] self.detect_multivariate_drift_calls = [] self.get_drift_metrics_dataframe_calls = [] def detect_univariate_drift(self, **kwargs): """Record arguments and return ``{}`` unless ``_univariate_side_effect`` overrides it.""" self.detect_univariate_drift_calls.append(kwargs) side_effect = type(self)._univariate_side_effect if isinstance(side_effect, Exception): raise side_effect if callable(side_effect): return side_effect(**kwargs) return {} def detect_multivariate_drift(self, **kwargs): """Record arguments and return ``{}`` unless ``_multivariate_side_effect`` overrides it.""" self.detect_multivariate_drift_calls.append(kwargs) side_effect = type(self)._multivariate_side_effect if isinstance(side_effect, Exception): raise side_effect if callable(side_effect): return side_effect(**kwargs) return {} def get_drift_metrics_dataframe(self, univariate_drift, multivariate_drift): """Return the configured drift dataframe (copy) or raise the configured exception.""" self.get_drift_metrics_dataframe_calls.append( {'univariate_drift': univariate_drift, 'multivariate_drift': multivariate_drift} ) if type(self)._drift_metrics_dataframe_exception is not None: raise type(self)._drift_metrics_dataframe_exception factory = type(self)._drift_response_factory if factory is None: return pd.DataFrame() result = factory(univariate_drift, multivariate_drift) return result.copy() if isinstance(result, pd.DataFrame) else result @classmethod def reset(cls): """Clear all configured side effects and the last constructed instance.""" cls._drift_response_factory = None cls._univariate_side_effect = None cls._multivariate_side_effect = None cls._drift_metrics_dataframe_exception = None cls.last_instance = None # Patch ``sientia_model.analytics.model_analysis.ModelAnalysis`` BEFORE the # production import chain runs so drift E2E tests can drive deterministic # analyzer outputs (the real implementation is data-dependent and would yield # values that drift across runs). The patch is applied to the real installed # module so that ``from laborious.activities.activities import Activities`` # resolves ``ModelAnalysis`` to ``_FakeModelAnalysis`` at import time. import sientia_model.analytics.model_analysis as _sientia_model_analysis_module # noqa: E402 _sientia_model_analysis_module.ModelAnalysis = _FakeModelAnalysis from laborious.activities.activities import Activities # noqa: E402 from laborious.workflows.drift import Drift from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.simple_metrics import SimpleMetrics from laborious.workflows.sub_workflows.format_and_export_prediction import ( FormatAndExportPrediction, ) from laborious.workflows.sub_workflows.prediction_process import PredictionProcess from sientia_do.notifications.handlers import CoreNotificationHandler from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController @pytest_asyncio.fixture(scope='session') def postgres_container(): """PostgreSQL testcontainer used by all E2E tests.""" postgres = PostgresContainer('postgres:15') postgres.start() yield postgres postgres.stop() @pytest_asyncio.fixture(scope='session') def minio_container(): """MinIO testcontainer used by E2E offload and payload retrieval paths.""" minio = MinioContainer() minio.start() yield minio minio.stop() @pytest_asyncio.fixture(scope='session') def mongo_container(): """MongoDB testcontainer used by real CoreNotificationHandler.""" mongo = DockerContainer('mongo:7').with_exposed_ports(27017) mongo.start() yield mongo mongo.stop() @pytest_asyncio.fixture def postgres_engine(postgres_container): """SQLAlchemy engine bound to the PostgreSQL testcontainer.""" engine = create_engine(postgres_container.get_connection_url()) yield engine engine.dispose() def _create_schema_and_tables(engine): """Create all schemas/tables required by workflow and activity paths.""" with engine.begin() as conn: conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema')) conn.execute( text( """ 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 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 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 TABLE IF NOT EXISTS predictions_schema.drift ( id SERIAL NOT NULL, model_id int4 NOT NULL, feature text NOT NULL, method text NOT NULL, value numeric NULL, drift bool NOT NULL, chunk int4 NOT NULL, "timestamp" timestamptz NOT NULL, timestamp_end text NULL, accurate bool NOT NULL, created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL, updated_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL, PRIMARY KEY (id) ); """ ) ) conn.execute( text( """ CREATE TABLE IF NOT EXISTS predictions_schema.simple_metrics_data ( id SERIAL NOT NULL, model_id int4 NOT NULL, metric text NOT NULL, value numeric NOT NULL, "timestamp" timestamptz NOT NULL, data_size int4 NOT NULL, interval_minutes int4 NOT NULL, created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL, PRIMARY KEY (id) ); """ ) ) conn.execute( text( """ CREATE TABLE IF NOT EXISTS predictions_schema.retrain_reports ( id SERIAL NOT NULL, model_id int4 NOT NULL, model_name text NOT NULL, "timestamp" text NOT NULL, status text NOT NULL, version text NULL, mlflow_run_id text NULL, mlflow_experiment_id text NULL, created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL, PRIMARY KEY (id) ); """ ) ) @pytest_asyncio.fixture(autouse=True) def setup_postgres_schema_and_tables(postgres_engine): """Ensure required schema and tables exist before each E2E test.""" _create_schema_and_tables(postgres_engine) yield @pytest.fixture(autouse=True) def reset_model_analysis_stub(): """Reset class-level state on the ModelAnalysis stub between tests.""" _FakeModelAnalysis.reset() yield _FakeModelAnalysis.reset() @pytest_asyncio.fixture def mock_logger(): """Logger double with readable console output for E2E runs.""" logger = MagicMock(spec=Logger) logger.info = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}')) logger.debug = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}')) logger.error = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}')) logger.warning = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}')) logger.custom_info = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}')) logger.custom_debug = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}')) logger.custom_error = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}')) logger.custom_warning = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}')) return logger @pytest_asyncio.fixture def metrics_controller(mock_logger): """Real metrics controller for E2E observability paths.""" return MetricsController(logger=mock_logger) @pytest_asyncio.fixture def notification_handler(mock_logger, mongo_container): """Real notification handler using MongoDB testcontainer.""" mongo_port = mongo_container.get_exposed_port(27017) handler = CoreNotificationHandler( connection_string=f'mongodb://localhost:{mongo_port}', database='test_db', logger=mock_logger, project_name='laborious', ) try: yield handler finally: handler.shutdown() @pytest.fixture def notification_inserts(notification_handler): """Spy on real Mongo insert calls issued by notification handler.""" collection = notification_handler.mongo_collection original_insert_one = collection.insert_one spy = MagicMock(wraps=original_insert_one) collection.insert_one = spy try: yield spy finally: collection.insert_one = original_insert_one class _FakeModelWrapper: """External MLflow wrapper double used by repository stub.""" def __init__(self): self.transform = MagicMock(side_effect=self._default_transform) self.predict = MagicMock(side_effect=self._default_predict) @staticmethod def _default_transform(data: pd.DataFrame): result = pd.DataFrame( { 'feature_1': [0.234] * len(data), 'feature_2': [0.783] * len(data), } ) result.index = data.index return result, {} @staticmethod def _default_predict(_params: dict, data: pd.DataFrame): pred = pd.DataFrame([0.5] * len(data), columns=['placeholder']) pred.index = data.index return pred, {} @pytest_asyncio.fixture def mlflow_repository_stub(): """External MLflow repository stub.""" repo = MagicMock() wrapper = _FakeModelWrapper() repo.stub_wrapper = wrapper repo.get_cached_model = MagicMock(return_value=wrapper) repo._client = MagicMock() return repo class _FakePIWebAPIClient: """External PI Web API client stub with deterministic responses.""" def __init__(self): self._responses = None self.write_value = MagicMock(side_effect=self._write_value) self.close = MagicMock() def set_side_effect(self, side_effect): self._responses = side_effect def _write_value(self, web_ids, value, metadata=None, **kwargs): if isinstance(self._responses, Exception): raise self._responses if isinstance(self._responses, list): item = self._responses.pop(0) if isinstance(item, Exception): raise item return item if callable(self._responses): return self._responses(web_ids=web_ids, value=value, metadata=metadata, **kwargs) return [{'WebId': wid, 'Errors': []} for wid in web_ids] @pytest_asyncio.fixture def pi_web_api_client_stub(): """PI Web API stub fixture.""" return _FakePIWebAPIClient() @pytest_asyncio.fixture def opc_repository_stub(): """OPC external dependency stub.""" repo = MagicMock() repo.write_data = MagicMock(return_value=(True, {'response_time': 0.1})) repo.disconnect = MagicMock() return repo @pytest_asyncio.fixture def plugin_store_stub(): """Plugin store external dependency stub.""" return MagicMock() @pytest_asyncio.fixture(scope='function') async def test_activities( postgres_container, minio_container, mock_logger, notification_handler, metrics_controller, mlflow_repository_stub, plugin_store_stub, pi_web_api_client_stub, opc_repository_stub, ): """Activities with real infra and external-system stubs only.""" 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': int(postgres_container.get_exposed_port(5432)), 'user': postgres_container.username, 'password': postgres_container.password, 'dbname': postgres_container.dbname, 'min_connections': 1, 'max_connections': 5, }, plugin_store=plugin_store_stub, 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, metrics_controller=metrics_controller, mlflow_repository=mlflow_repository_stub, ) activities.pi_web_api_client = pi_web_api_client_stub activities.opc_repository = {'1': opc_repository_stub} try: yield activities finally: activities.shutdown() @pytest_asyncio.fixture(scope='function') async def test_activities_real_minio( postgres_container, minio_container, mock_logger, notification_handler, metrics_controller, mlflow_repository_stub, plugin_store_stub, pi_web_api_client_stub, opc_repository_stub, ): """Compatibility alias 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': int(postgres_container.get_exposed_port(5432)), 'user': postgres_container.username, 'password': postgres_container.password, 'dbname': postgres_container.dbname, 'min_connections': 1, 'max_connections': 5, }, plugin_store=plugin_store_stub, 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, metrics_controller=metrics_controller, mlflow_repository=mlflow_repository_stub, ) activities.pi_web_api_client = pi_web_api_client_stub activities.opc_repository = {'1': opc_repository_stub} try: yield activities finally: activities.shutdown() def _worker_activity_list(test_activities: Activities): """List of registered activity callables used by Temporal worker in E2E.""" return [ 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(): """Temporal test environment with time-skipping.""" 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): """Temporal worker for full predictions-batch and child workflows.""" with ThreadPoolExecutor(max_workers=32) as activity_executor: async with Worker( temporal_test_env.client, task_queue='test-queue', workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], activities=_worker_activity_list(test_activities), activity_executor=activity_executor, ) as worker: yield worker @pytest_asyncio.fixture(scope='function') async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio): """Temporal worker alias for tests that emphasize MinIO behavior.""" with ThreadPoolExecutor(max_workers=32) as activity_executor: async with Worker( temporal_test_env.client, task_queue='test-queue', workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], activities=_worker_activity_list(test_activities_real_minio), activity_executor=activity_executor, ) as worker: yield worker def _drift_worker_activity_list(test_activities: Activities): """Activity callables registered on the drift Temporal worker.""" return [ test_activities.load_custom_query, test_activities.get_reference_data, test_activities.calculate_drift, test_activities.export_data_to_postgres, ] @pytest_asyncio.fixture(scope='function') async def temporal_worker_drift(temporal_test_env, test_activities): """Temporal worker registered with the Drift workflow and its activities.""" with ThreadPoolExecutor(max_workers=32) as activity_executor: async with Worker( temporal_test_env.client, task_queue='test-queue', workflows=[Drift], activities=_drift_worker_activity_list(test_activities), activity_executor=activity_executor, ) as worker: yield worker def _simple_metrics_worker_activity_list(test_activities: Activities): """Activity callables registered on the simple-metrics Temporal worker.""" return [ test_activities.load_custom_query, test_activities.calculate_simple_metrics, test_activities.export_data_to_postgres, ] @pytest_asyncio.fixture(scope='function') async def temporal_worker_simple_metrics(temporal_test_env, test_activities): """Temporal worker registered with the SimpleMetrics workflow and its activities.""" with ThreadPoolExecutor(max_workers=32) as activity_executor: async with Worker( temporal_test_env.client, task_queue='test-queue', workflows=[SimpleMetrics], activities=_simple_metrics_worker_activity_list(test_activities), activity_executor=activity_executor, ) as worker: yield worker def _minimal_retrain_worker_activity_list(test_activities: Activities): """Activity callables registered on the minimal-retrain Temporal worker.""" return [ test_activities.load_query_with_minio_offload, test_activities.retrain_model, test_activities.update_production_model, test_activities.format_retrain_report, test_activities.export_data_to_postgres, ] @pytest_asyncio.fixture(scope='function') async def temporal_worker_minimal_retrain(temporal_test_env, test_activities): """Temporal worker registered with the MinimalRetrain workflow and its activities.""" with ThreadPoolExecutor(max_workers=32) as activity_executor: async with Worker( temporal_test_env.client, task_queue='test-queue', workflows=[MinimalRetrain], activities=_minimal_retrain_worker_activity_list(test_activities), activity_executor=activity_executor, ) as worker: yield worker @pytest.fixture def model_analysis_stub(): """ Expose the ``_FakeModelAnalysis`` class so drift tests can configure responses. Use class-level attributes to control the analyzer outputs: - ``stub.set_drift_dataframe(factory)`` to provide rows for ``calculate_drift``. - ``stub.raise_on_get_drift_metrics_dataframe(exception)`` to simulate failures. """ class _Helper: """Thin convenience wrapper around _FakeModelAnalysis class state.""" cls = _FakeModelAnalysis def set_drift_dataframe(self, factory): self.cls._drift_response_factory = factory def raise_on_get_drift_metrics_dataframe(self, exc: Exception): self.cls._drift_metrics_dataframe_exception = exc def set_univariate_side_effect(self, side_effect): self.cls._univariate_side_effect = side_effect def set_multivariate_side_effect(self, side_effect): self.cls._multivariate_side_effect = side_effect @property def last_instance(self): return self.cls.last_instance return _Helper()