- Added a new fixture to manage runtime report artifacts in a writable temp directory during E2E tests, addressing permission issues in local CI/dev environments. - Updated `conftest.py` to include a requirements.txt file in the model packaging path for training activities. - Refactored existing fixtures to use `pytest.fixture` instead of `pytest_asyncio.fixture` for better compatibility. - Enhanced the `Reports` class to include a target alias for report metrics, ensuring compatibility with Evidently's reporting requirements. - Introduced new test scenarios to validate the handling of missing and whitespace-only `date_column` inputs in the training workflow. These changes improve the robustness of the E2E testing framework and enhance the clarity of model reporting metrics.
698 lines
22 KiB
Python
698 lines
22 KiB
Python
"""
|
||
Pytest configuration and fixtures for E2E tests.
|
||
|
||
External dependencies use testcontainers or real SDK integrations (no mocks of
|
||
model_manager or other first-party code):
|
||
|
||
- PostgreSQL, MinIO, MongoDB, Gitea: testcontainers.
|
||
- MLflow: real client with ``file://`` tracking URI (no MLflow server process).
|
||
- Temporal: ``WorkflowEnvironment.start_time_skipping()`` — official in-process
|
||
test runtime from temporalio; exercises real workflows and activity code, not
|
||
stubs of business logic.
|
||
- Observability: ``Logger`` (``get_logger`` from ``model_manager.utils.logger_helper``)
|
||
and ``MetricsController`` from sientia_do, same stack as production.
|
||
"""
|
||
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
import base64
|
||
import csv
|
||
import io
|
||
import os
|
||
import shutil
|
||
import tempfile
|
||
import time
|
||
|
||
import mlflow
|
||
import pytest
|
||
import pytest_asyncio
|
||
import requests
|
||
from minio import Minio
|
||
from sqlalchemy import create_engine, text
|
||
from testcontainers.core.container import DockerContainer # type: ignore[import-untyped]
|
||
from testcontainers.minio import MinioContainer # type: ignore[import-untyped]
|
||
from testcontainers.mongodb import MongoDbContainer # type: ignore[import-untyped]
|
||
from testcontainers.postgres import PostgresContainer # type: ignore[import-untyped]
|
||
from temporalio.testing import WorkflowEnvironment
|
||
from temporalio.worker import Worker
|
||
|
||
from model_manager.activities.activities import Activities
|
||
from model_manager.utils.logger_helper import get_logger
|
||
from model_manager.workflows.cleanup_files import CleanupFiles
|
||
from model_manager.workflows.train_model import TrainModel
|
||
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||
from sientia_do.observability.metrics_controller import MetricsController
|
||
from sientia_model.model_repository.plugin_store import PluginStore
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CSV training data: columns must match the variable_columns and target_variable
|
||
# used across all test scenarios.
|
||
_TRAIN_CSV_COLUMNS = [
|
||
'timestamp',
|
||
'303-WIT-200(Value)',
|
||
'03CV020/CORRENTE_N_M1_PV(Value)',
|
||
'303-WIT-230(Value)',
|
||
'03CV022/CORRENTE_N_M1_PV(Value)',
|
||
]
|
||
_MINIO_BUCKET = 'model-training'
|
||
_MINIO_OBJECT = 'training_data.csv'
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers – CSV generation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _build_training_csv() -> bytes:
|
||
"""
|
||
Generate a 150-row CSV with all columns needed by test scenarios.
|
||
|
||
The numeric values cycle deterministically so lags and static-window
|
||
removal always find enough rows in both train and validation splits.
|
||
"""
|
||
output = io.StringIO()
|
||
writer = csv.writer(output)
|
||
writer.writerow(_TRAIN_CSV_COLUMNS)
|
||
for i in range(150):
|
||
ts = f'2025-06-{(i // 24) + 2:02d} {i % 24:02d}:00:00'
|
||
wit200 = round(30.0 + (i % 20) * 0.5, 2)
|
||
cv020 = round(100.0 + (i % 15) * 0.3, 2)
|
||
wit230 = round(25.0 + (i % 18) * 0.4, 2)
|
||
cv022 = round(90.0 + (i % 12) * 0.25, 2)
|
||
writer.writerow([ts, wit200, cv020, wit230, cv022])
|
||
return output.getvalue().encode('utf-8')
|
||
|
||
|
||
def _build_training_csv_dd_mm_yyyy() -> bytes:
|
||
"""
|
||
Generate a 150-row CSV with dd/MM/yyyy HH:mm:ss timestamps and
|
||
a DATA column header, for scenarios 12/13 that use a different date format.
|
||
"""
|
||
output = io.StringIO()
|
||
writer = csv.writer(output)
|
||
writer.writerow([
|
||
'DATA',
|
||
'303-WIT-230(Value)',
|
||
'03CV022/CORRENTE_N_M1_PV(Value)',
|
||
])
|
||
for i in range(150):
|
||
day = (i % 30) + 1
|
||
ts = f'{day:02d}/05/2022 {i % 24:02d}:00:00'
|
||
wit230 = round(25.0 + (i % 18) * 0.4, 2)
|
||
cv022 = round(90.0 + (i % 12) * 0.25, 2)
|
||
writer.writerow([ts, wit230, cv022])
|
||
return output.getvalue().encode('utf-8')
|
||
|
||
|
||
def _build_training_csv_custom_target_column() -> bytes:
|
||
"""
|
||
Same layout as the standard CSV but the target column has a non-default name
|
||
(not ``target``) to exercise report and metrics paths.
|
||
"""
|
||
output = io.StringIO()
|
||
writer = csv.writer(output)
|
||
writer.writerow([
|
||
'timestamp',
|
||
'303-WIT-200(Value)',
|
||
'MY_CUSTOM_TARGET_COLUMN',
|
||
])
|
||
for i in range(150):
|
||
ts = f'2025-06-{(i // 24) + 2:02d} {i % 24:02d}:00:00'
|
||
wit200 = round(30.0 + (i % 20) * 0.5, 2)
|
||
target_val = round(100.0 + (i % 15) * 0.3, 2)
|
||
writer.writerow([ts, wit200, target_val])
|
||
return output.getvalue().encode('utf-8')
|
||
|
||
|
||
def _build_training_csv_timestamp_header_naive() -> bytes:
|
||
"""
|
||
Naive datetimes under column ``Timestamp`` (common UI export) for scenario 16.
|
||
"""
|
||
output = io.StringIO()
|
||
writer = csv.writer(output)
|
||
writer.writerow([
|
||
'Timestamp',
|
||
'303-WIT-200(Value)',
|
||
'03CV020/CORRENTE_N_M1_PV(Value)',
|
||
])
|
||
for i in range(150):
|
||
ts = f'2025-06-{(i // 24) + 2:02d} {i % 24:02d}:00:00'
|
||
wit200 = round(30.0 + (i % 20) * 0.5, 2)
|
||
cv020 = round(100.0 + (i % 15) * 0.3, 2)
|
||
writer.writerow([ts, wit200, cv020])
|
||
return output.getvalue().encode('utf-8')
|
||
|
||
|
||
def _build_training_csv_blank_timestamp_row() -> bytes:
|
||
"""Standard columns with one row where ``timestamp`` is empty (NaN after parse)."""
|
||
output = io.StringIO()
|
||
writer = csv.writer(output)
|
||
writer.writerow(_TRAIN_CSV_COLUMNS)
|
||
for i in range(150):
|
||
wit200 = round(30.0 + (i % 20) * 0.5, 2)
|
||
cv020 = round(100.0 + (i % 15) * 0.3, 2)
|
||
wit230 = round(25.0 + (i % 18) * 0.4, 2)
|
||
cv022 = round(90.0 + (i % 12) * 0.25, 2)
|
||
if i == 17:
|
||
writer.writerow(['', wit200, cv020, wit230, cv022])
|
||
else:
|
||
ts = f'2025-06-{(i // 24) + 2:02d} {i % 24:02d}:00:00'
|
||
writer.writerow([ts, wit200, cv020, wit230, cv022])
|
||
return output.getvalue().encode('utf-8')
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers – Gitea seed
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _wait_for_gitea(base_url: str, timeout: int = 120) -> None:
|
||
"""Poll Gitea until it responds to HTTP requests."""
|
||
deadline = time.time() + timeout
|
||
last_err = None
|
||
while time.time() < deadline:
|
||
try:
|
||
resp = requests.get(f'{base_url}/', timeout=3)
|
||
if resp.status_code in (200, 404, 302):
|
||
return
|
||
except Exception as e:
|
||
last_err = e
|
||
time.sleep(2)
|
||
raise TimeoutError(f'Gitea did not start within {timeout}s at {base_url}. Last error: {last_err}')
|
||
|
||
|
||
def _gitea_api(method: str, url: str, auth: tuple, **kwargs) -> requests.Response:
|
||
resp = requests.request(method, url, auth=auth, timeout=30, **kwargs)
|
||
try:
|
||
resp.raise_for_status()
|
||
except requests.exceptions.HTTPError as e:
|
||
raise RuntimeError(f"Gitea API error {resp.status_code}: {resp.text}") from e
|
||
return resp
|
||
|
||
|
||
def _seed_gitea(base_url: str, admin_user: str, admin_pass: str) -> None:
|
||
"""
|
||
Create a fictitious model-store repository with dummy models.
|
||
"""
|
||
auth = (admin_user, admin_pass)
|
||
api = f'{base_url}/api/v1'
|
||
|
||
# Create repository
|
||
_gitea_api(
|
||
'POST', f'{api}/user/repos', auth,
|
||
json={'name': 'model-store', 'private': False, 'auto_init': False},
|
||
)
|
||
|
||
# Root index.yaml
|
||
root_index = """
|
||
store_name: "E2E Test Store"
|
||
version: 1
|
||
models:
|
||
- name: "linear_regression"
|
||
version: 1
|
||
runtime: "basic"
|
||
- name: "polynomial_regression"
|
||
version: 1
|
||
runtime: "basic"
|
||
runtimes:
|
||
basic:
|
||
version: "1.0.0"
|
||
libraries:
|
||
- name: "pandas"
|
||
- name: "numpy"
|
||
"""
|
||
|
||
# Model index.yaml (shared for all dummies)
|
||
model_index = """
|
||
name: "{model_name}"
|
||
version: 1
|
||
runtime: "basic"
|
||
path: "wrapper.py"
|
||
class: "DummyWrapper"
|
||
model:
|
||
class: "DummyModel"
|
||
path: "model_logic.py"
|
||
external: false
|
||
data_model:
|
||
class: "DummyTransformer"
|
||
path: "model_logic.py"
|
||
external: false
|
||
"""
|
||
|
||
# schemas.yaml
|
||
schemas_yaml = """
|
||
model:
|
||
type: object
|
||
properties: {}
|
||
data_model:
|
||
type: object
|
||
properties: {}
|
||
opt_params:
|
||
type: object
|
||
properties: {}
|
||
"""
|
||
|
||
# wrapper.py
|
||
wrapper_py = """
|
||
from sientia_model.wrappers.sientia_model import SientiaModel
|
||
import pandas as pd
|
||
import numpy as np
|
||
from typing import Any
|
||
|
||
class DummyWrapper(SientiaModel):
|
||
def _predict(self, data: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, Any]]:
|
||
self._log("info", f"Predicting dummy model for {self.model_type}")
|
||
# Return a simple prediction (mean or 0.5) to allow metrics computation
|
||
preds = pd.DataFrame({self.target: [0.5] * len(data)}, index=data.index)
|
||
return preds, {}
|
||
|
||
def _transform(self, data: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, Any]]:
|
||
return data, {}
|
||
|
||
def _train_transformer(self, train_data: pd.DataFrame, val_data: pd.DataFrame) -> None:
|
||
pass
|
||
|
||
def _train_model(self, x: pd.DataFrame, y: pd.DataFrame, x_val: pd.DataFrame | None = None, y_val: pd.DataFrame | None = None) -> None:
|
||
self.target = y.columns[0]
|
||
|
||
def _retrain_transformer(self, data: pd.DataFrame) -> None:
|
||
pass
|
||
|
||
def _retrain_model(self, x: pd.DataFrame, y: pd.DataFrame | None) -> None:
|
||
pass
|
||
"""
|
||
|
||
# model_logic.py
|
||
model_logic_py = """
|
||
class DummyModel:
|
||
def __init__(self, **kwargs):
|
||
pass
|
||
|
||
class DummyTransformer:
|
||
def __init__(self, **kwargs):
|
||
pass
|
||
"""
|
||
|
||
# requirements.txt required by current model packaging path in training activity.
|
||
requirements_txt = """
|
||
pandas
|
||
numpy
|
||
"""
|
||
|
||
def push_file(path: str, content: str):
|
||
encoded = base64.b64encode(content.encode()).decode()
|
||
_gitea_api(
|
||
'POST',
|
||
f'{api}/repos/{admin_user}/model-store/contents/{path}',
|
||
auth,
|
||
json={'message': f'seed: {path}', 'content': encoded},
|
||
)
|
||
|
||
# Push root index
|
||
push_file('index.yaml', root_index)
|
||
|
||
# Push files for both models used in tests
|
||
for model_name in ['linear_regression', 'polynomial_regression']:
|
||
prefix = f'models/{model_name}'
|
||
push_file(f'{prefix}/index.yaml', model_index.format(model_name=model_name))
|
||
push_file(f'{prefix}/schemas.yaml', schemas_yaml)
|
||
push_file(f'{prefix}/wrapper.py', wrapper_py)
|
||
push_file(f'{prefix}/model_logic.py', model_logic_py)
|
||
push_file(f'{prefix}/requirements.txt', requirements_txt.strip() + '\n')
|
||
push_file(f'{prefix}/__init__.py', "")
|
||
|
||
# Push runtime
|
||
push_file('runtime/basic.yaml', 'name: basic\nversion: "1.0.0"\nlibraries: []')
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Session-scoped containers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.fixture(scope='session')
|
||
def postgres_container():
|
||
"""PostgreSQL 15 container for experiment_run table."""
|
||
container = PostgresContainer('postgres:15')
|
||
container.start()
|
||
yield container
|
||
container.stop()
|
||
|
||
|
||
@pytest.fixture(scope='session')
|
||
def minio_container():
|
||
"""MinIO container for training CSV storage."""
|
||
container = MinioContainer()
|
||
container.start()
|
||
yield container
|
||
container.stop()
|
||
|
||
|
||
@pytest.fixture(scope='session')
|
||
def mongodb_container():
|
||
"""MongoDB container for CoreNotificationHandler."""
|
||
container = MongoDbContainer('mongo:7')
|
||
container.start()
|
||
yield container
|
||
container.stop()
|
||
|
||
|
||
@pytest.fixture(scope='session')
|
||
def gitea_container():
|
||
"""
|
||
Gitea container with a ``model-store`` repo seeded via REST API
|
||
(``index.yaml``, dummy model files, ``_seed_gitea``).
|
||
|
||
The container starts with INSTALL_LOCK so no setup wizard is needed.
|
||
An admin user is created via Gitea's CLI before the HTTP API is used.
|
||
"""
|
||
admin_user = 'gitea_admin'
|
||
admin_pass = 'gitea_admin_pass' # noqa: S105
|
||
|
||
container = (
|
||
DockerContainer('gitea/gitea:latest')
|
||
.with_env('GITEA__security__INSTALL_LOCK', 'true')
|
||
.with_env('GITEA__server__HTTP_PORT', '3000')
|
||
.with_env('GITEA__log__LEVEL', 'Warn')
|
||
.with_exposed_ports(3000)
|
||
)
|
||
container.start()
|
||
|
||
port = container.get_exposed_port(3000)
|
||
base_url = f'http://localhost:{port}'
|
||
|
||
_wait_for_gitea(base_url)
|
||
time.sleep(5) # Wait a bit for DB to fully initialize after HTTP is up
|
||
|
||
# Create admin user via Gitea CLI inside the container
|
||
# Must run after Gitea is fully initialized
|
||
gitea_cmd = (
|
||
f'gitea admin user create '
|
||
f'--username {admin_user} '
|
||
f'--password {admin_pass} '
|
||
f'--email admin@test.local '
|
||
f'--admin '
|
||
f'--must-change-password=false'
|
||
)
|
||
exec_result = container.exec(f"su git -c '{gitea_cmd}'")
|
||
if exec_result.exit_code != 0:
|
||
raise RuntimeError(f"Failed to create Gitea admin user: {exec_result.output.decode('utf-8')}")
|
||
|
||
_seed_gitea(base_url, admin_user, admin_pass)
|
||
|
||
yield {
|
||
'container': container,
|
||
'base_url': base_url,
|
||
'admin_user': admin_user,
|
||
'admin_pass': admin_pass,
|
||
}
|
||
|
||
container.stop()
|
||
|
||
|
||
@pytest.fixture(scope='session')
|
||
def mlflow_tracking_dir():
|
||
"""Local MLflow filesystem tracking directory (no network needed)."""
|
||
tmpdir = tempfile.mkdtemp(prefix='mlflow-e2e-')
|
||
mlflow.set_tracking_uri(f'file://{tmpdir}')
|
||
yield tmpdir
|
||
shutil.rmtree(tmpdir, ignore_errors=True)
|
||
|
||
|
||
@pytest.fixture(scope='session', autouse=True)
|
||
def e2e_runtime_reports_dir():
|
||
"""
|
||
Route runtime report artifacts to a writable temp directory during E2E.
|
||
|
||
Production defaults point to /var/lib/model-manager; in local CI/dev runs this
|
||
path may be unavailable. This fixture keeps the same code paths while avoiding
|
||
host permission issues.
|
||
"""
|
||
import model_manager.runtime_paths as runtime_paths
|
||
import model_manager.utils.repository.data_manager_repository as data_repo_module
|
||
|
||
base_dir = tempfile.mkdtemp(prefix='model-manager-e2e-runtime-')
|
||
reports_root = f'{base_dir}/reports'
|
||
reports_temp_dir = f'{reports_root}/temp'
|
||
project_base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'model_manager'))
|
||
|
||
old_runtime_reports_root = runtime_paths.REPORTS_ROOT
|
||
old_runtime_reports_temp = runtime_paths.REPORTS_TEMP_DIR
|
||
old_runtime_project_base = runtime_paths.PROJECT_BASE_PATH
|
||
old_repo_reports_root = data_repo_module.REPORTS_ROOT
|
||
old_repo_project_base = data_repo_module.PROJECT_BASE_PATH
|
||
|
||
runtime_paths.REPORTS_ROOT = reports_root
|
||
runtime_paths.REPORTS_TEMP_DIR = reports_temp_dir
|
||
runtime_paths.PROJECT_BASE_PATH = project_base_path
|
||
data_repo_module.REPORTS_ROOT = reports_root
|
||
data_repo_module.PROJECT_BASE_PATH = project_base_path
|
||
|
||
try:
|
||
yield reports_root
|
||
finally:
|
||
runtime_paths.REPORTS_ROOT = old_runtime_reports_root
|
||
runtime_paths.REPORTS_TEMP_DIR = old_runtime_reports_temp
|
||
runtime_paths.PROJECT_BASE_PATH = old_runtime_project_base
|
||
data_repo_module.REPORTS_ROOT = old_repo_reports_root
|
||
data_repo_module.PROJECT_BASE_PATH = old_repo_project_base
|
||
shutil.rmtree(base_dir, ignore_errors=True)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Session-scoped: seed MinIO with training CSV
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.fixture(scope='session', autouse=True)
|
||
def upload_training_csv(minio_container, mlflow_tracking_dir): # noqa: ARG001
|
||
"""
|
||
Upload training CSV files to the MinIO container before any test runs.
|
||
Depends on mlflow_tracking_dir to ensure the MLflow URI is set at session start.
|
||
"""
|
||
port = minio_container.get_exposed_port(9000)
|
||
client = Minio(
|
||
f'localhost:{port}',
|
||
access_key='minioadmin',
|
||
secret_key='minioadmin',
|
||
secure=False,
|
||
)
|
||
|
||
if not client.bucket_exists(_MINIO_BUCKET):
|
||
client.make_bucket(_MINIO_BUCKET)
|
||
|
||
# Standard training CSV
|
||
csv_bytes = _build_training_csv()
|
||
client.put_object(
|
||
_MINIO_BUCKET,
|
||
_MINIO_OBJECT,
|
||
io.BytesIO(csv_bytes),
|
||
length=len(csv_bytes),
|
||
content_type='text/csv',
|
||
)
|
||
|
||
# dd/MM/yyyy format CSV for scenarios 12/13
|
||
alt_csv_bytes = _build_training_csv_dd_mm_yyyy()
|
||
client.put_object(
|
||
_MINIO_BUCKET,
|
||
'training_data_dd_mm_yyyy.csv',
|
||
io.BytesIO(alt_csv_bytes),
|
||
length=len(alt_csv_bytes),
|
||
content_type='text/csv',
|
||
)
|
||
|
||
custom_target = _build_training_csv_custom_target_column()
|
||
client.put_object(
|
||
_MINIO_BUCKET,
|
||
'training_data_custom_target.csv',
|
||
io.BytesIO(custom_target),
|
||
length=len(custom_target),
|
||
content_type='text/csv',
|
||
)
|
||
|
||
ts_header = _build_training_csv_timestamp_header_naive()
|
||
client.put_object(
|
||
_MINIO_BUCKET,
|
||
'training_data_timestamp_naive.csv',
|
||
io.BytesIO(ts_header),
|
||
length=len(ts_header),
|
||
content_type='text/csv',
|
||
)
|
||
|
||
blank_ts = _build_training_csv_blank_timestamp_row()
|
||
client.put_object(
|
||
_MINIO_BUCKET,
|
||
'training_data_blank_timestamp_row.csv',
|
||
io.BytesIO(blank_ts),
|
||
length=len(blank_ts),
|
||
content_type='text/csv',
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Function-scoped: database engine + schema setup
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.fixture
|
||
def postgres_engine(postgres_container):
|
||
"""SQLAlchemy engine connected to the test PostgreSQL container."""
|
||
engine = create_engine(postgres_container.get_connection_url())
|
||
yield engine
|
||
engine.dispose()
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def setup_experiment_run_table(postgres_engine):
|
||
"""
|
||
Create the experiment_run table before each test and drop it afterwards
|
||
to guarantee full isolation between tests.
|
||
"""
|
||
with postgres_engine.begin() as conn:
|
||
conn.execute(text("""
|
||
CREATE TABLE IF NOT EXISTS public.experiment_run (
|
||
id INT PRIMARY KEY,
|
||
experiment_name TEXT NOT NULL,
|
||
run_name TEXT,
|
||
username TEXT,
|
||
status TEXT NOT NULL DEFAULT 'ORCHESTRATOR_WAITING_PROC',
|
||
error_message TEXT,
|
||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||
bucket_name TEXT,
|
||
file_name TEXT
|
||
)
|
||
"""))
|
||
yield
|
||
with postgres_engine.begin() as conn:
|
||
conn.execute(text('DROP TABLE IF EXISTS public.experiment_run'))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Observability (real sientia_do implementations)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.fixture(scope='session')
|
||
def e2e_logger():
|
||
"""Shared production-style Logger for the whole E2E session."""
|
||
return get_logger('model-manager-e2e')
|
||
|
||
|
||
@pytest.fixture
|
||
def metrics_controller(e2e_logger):
|
||
"""MetricsController bound to the E2E logger (fresh instance per test)."""
|
||
return MetricsController(logger=e2e_logger)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Application fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.fixture
|
||
def notification_handler(mongodb_container, e2e_logger):
|
||
"""
|
||
Real CoreNotificationHandler connected to the MongoDB testcontainer.
|
||
"""
|
||
connection_url = mongodb_container.get_connection_url()
|
||
handler = CoreNotificationHandler(
|
||
connection_string=connection_url,
|
||
database='test_notifications',
|
||
logger=e2e_logger,
|
||
project_name='model-manager-e2e',
|
||
)
|
||
yield handler
|
||
handler.shutdown()
|
||
|
||
|
||
@pytest.fixture
|
||
def plugin_store(gitea_container, e2e_logger, metrics_controller, notification_handler):
|
||
"""
|
||
Real PluginStore pointed at the Gitea testcontainer.
|
||
cache_ttl_seconds=0 forces a fresh download every test.
|
||
"""
|
||
store = PluginStore(
|
||
base_url=gitea_container['base_url'],
|
||
owner=gitea_container['admin_user'],
|
||
repo='model-store',
|
||
username=gitea_container['admin_user'],
|
||
password=gitea_container['admin_pass'],
|
||
cache_ttl_seconds=0,
|
||
logger=e2e_logger,
|
||
notification_handler=notification_handler,
|
||
metrics_controller=metrics_controller,
|
||
)
|
||
yield store
|
||
|
||
|
||
@pytest.fixture
|
||
def test_activities(
|
||
postgres_container,
|
||
minio_container,
|
||
mlflow_tracking_dir, # noqa: ARG001 – ensures MLflow URI is set
|
||
plugin_store,
|
||
e2e_logger,
|
||
notification_handler,
|
||
metrics_controller,
|
||
):
|
||
"""
|
||
Real Activities instance wired to all testcontainers.
|
||
"""
|
||
pg_port = postgres_container.get_exposed_port(5432)
|
||
minio_port = minio_container.get_exposed_port(9000)
|
||
|
||
activities = Activities(
|
||
postgres_config={
|
||
'host': 'localhost',
|
||
'port': int(pg_port),
|
||
'user': 'test',
|
||
'password': 'test',
|
||
'dbname': 'test',
|
||
'min_connections': 1,
|
||
'max_connections': 5,
|
||
},
|
||
mlflow_config={
|
||
'url': mlflow.get_tracking_uri(),
|
||
'username': None,
|
||
'password': None,
|
||
},
|
||
minio_config={
|
||
'endpoint_url': f'http://localhost:{minio_port}',
|
||
'access_key': 'minioadmin',
|
||
'secret_key': 'minioadmin',
|
||
'use_ssl': False,
|
||
'default_bucket': _MINIO_BUCKET,
|
||
},
|
||
plugin_store=plugin_store,
|
||
logger=e2e_logger,
|
||
notification_handler=notification_handler,
|
||
metrics_controller=metrics_controller,
|
||
)
|
||
yield activities
|
||
activities.shutdown()
|
||
|
||
|
||
def _activity_list(activities: Activities) -> list:
|
||
return [
|
||
activities.update_experiment_run,
|
||
activities.load_model_metadata,
|
||
activities.validate_train_params,
|
||
activities.train_model,
|
||
activities.cleanup_resources,
|
||
activities.cleanup_temp_directories,
|
||
]
|
||
|
||
|
||
@pytest_asyncio.fixture(scope='function')
|
||
async def temporal_test_env():
|
||
"""Temporal SDK test environment (time-skipping); runs real workflow/activity code."""
|
||
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 registered with all workflows and activities."""
|
||
with ThreadPoolExecutor() as executor:
|
||
async with Worker(
|
||
temporal_test_env.client,
|
||
task_queue='test-queue',
|
||
workflows=[TrainModel, CleanupFiles],
|
||
activities=_activity_list(test_activities),
|
||
activity_executor=executor,
|
||
) as worker:
|
||
yield worker
|