10c7e292b9
Refactor ModelMetrics to utilize DriftAnalysis for drift detection - Replaced ModelAnalysis with DriftAnalysis in the ModelMetrics class to enhance drift detection capabilities. - Updated method signatures and documentation to reflect the changes in target_name and return values. - Adjusted data handling to ensure compatibility with the new analysis methods and improved clarity in the drift metrics dataframe preparation.
464 lines
16 KiB
Python
464 lines
16 KiB
Python
"""Pytest configuration and fixtures for E2E tests."""
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
import pytest_asyncio
|
|
from sqlalchemy import create_engine
|
|
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
|
|
|
|
from laborious.activities.activities import Activities
|
|
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
|
|
|
|
# Single source of truth for the test database schema. Mirrors the production
|
|
# DDL for ``sientia_data`` so any production change can be pasted directly into
|
|
# this file (see ``e2e/db_schema.sql``) without touching Python.
|
|
DB_SCHEMA_SQL_PATH = Path(__file__).parent / 'db_schema.sql'
|
|
|
|
|
|
@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.
|
|
|
|
Loads the DDL from ``e2e/db_schema.sql`` (single source of truth that
|
|
mirrors the production schema). The SQL file is executed via the raw
|
|
DBAPI cursor so multi-statement DDL is supported.
|
|
"""
|
|
sql_text = DB_SCHEMA_SQL_PATH.read_text(encoding='utf-8')
|
|
with engine.begin() as conn:
|
|
conn.exec_driver_sql(sql_text)
|
|
|
|
|
|
@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_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
|
|
|
|
|