e6018af23f
Update README, requirements, and E2E tests for improved configuration and functionality - Enhanced the README with updated model configuration examples, including the addition of an alias for production. - Removed the `requirements-light.txt` file and updated `requirements-local.txt` and `requirements.txt` to replace `asyncua` with `opcua`. - Refactored E2E test scenarios to utilize scenario input files for better maintainability and clarity. - Improved test coverage for MinIO offload functionality and added new helper functions for loading scenario inputs. - Updated `values.yaml` to reflect new global configurations and environment variables for the laborious worker.
686 lines
24 KiB
Python
686 lines
24 KiB
Python
"""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()
|