6bd30e3328
- Updated `pyproject.toml` to include new linting rules for end-to-end tests. - Modified `requirements-dev.txt` to add dependencies for E2E testing with `testcontainers` and `requests`. - Refactored multiple JSON test scenario files to standardize structure, including new fields for `experiment_run_id`, `bucket_name`, and `file_name`. - Improved model training parameters in `train_model_params.py` to use `experiment_name` directly. - Adjusted `data_manager_repository.py` to utilize the updated `experiment_name` for logging. These changes improve the organization and clarity of regression model tests and enhance the overall testing framework.
582 lines
18 KiB
Python
582 lines
18 KiB
Python
"""
|
||
Pytest configuration and fixtures for E2E tests.
|
||
|
||
All external dependencies use real services:
|
||
- PostgreSQL: testcontainers (postgres:15)
|
||
- MinIO: testcontainers (minio)
|
||
- MongoDB: testcontainers (mongo:7)
|
||
- MLflow: local filesystem tracking (no network)
|
||
- Gitea: testcontainers generic container (gitea/gitea:latest),
|
||
seeded with model-plugin-warehouse files via REST API
|
||
- Temporal: in-memory WorkflowEnvironment (time-skipping)
|
||
"""
|
||
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
import base64
|
||
import csv
|
||
import io
|
||
import os
|
||
import shutil
|
||
import tempfile
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
from unittest.mock import MagicMock
|
||
|
||
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
|
||
from testcontainers.minio import MinioContainer
|
||
from testcontainers.mongodb import MongoDbContainer
|
||
from testcontainers.postgres import PostgresContainer
|
||
from temporalio.testing import WorkflowEnvironment
|
||
from temporalio.worker import Worker
|
||
|
||
from model_manager.activities.activities import Activities
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Paths
|
||
# ---------------------------------------------------------------------------
|
||
_WAREHOUSE_ROOT = Path(
|
||
'/home/grezewave/Documents/projects/sientia/model-plugin-warehouse'
|
||
)
|
||
|
||
# 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')
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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
|
||
"""
|
||
|
||
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}/__init__.py', "")
|
||
|
||
# Push runtime
|
||
push_file('runtime/basic.yaml', 'name: basic\nversion: "1.0.0"\nlibraries: []')
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Session-scoped containers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.fixture(scope='session')
|
||
def postgres_container():
|
||
"""PostgreSQL 15 container for experiment_run table."""
|
||
container = PostgresContainer('postgres:15')
|
||
container.start()
|
||
yield container
|
||
container.stop()
|
||
|
||
|
||
@pytest_asyncio.fixture(scope='session')
|
||
def minio_container():
|
||
"""MinIO container for training CSV storage."""
|
||
container = MinioContainer()
|
||
container.start()
|
||
yield container
|
||
container.stop()
|
||
|
||
|
||
@pytest_asyncio.fixture(scope='session')
|
||
def mongodb_container():
|
||
"""MongoDB container for CoreNotificationHandler."""
|
||
container = MongoDbContainer('mongo:7')
|
||
container.start()
|
||
yield container
|
||
container.stop()
|
||
|
||
|
||
@pytest_asyncio.fixture(scope='session')
|
||
def gitea_container():
|
||
"""
|
||
Gitea container seeded with the model-plugin-warehouse files.
|
||
|
||
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)
|
||
import time
|
||
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_asyncio.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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Session-scoped: seed MinIO with training CSV
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.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',
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Function-scoped: database engine + schema setup
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.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_asyncio.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'))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Mock-only fixtures (no external service equivalent)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.fixture
|
||
def mock_logger():
|
||
"""Minimal logger that prints to stdout (no external observability needed)."""
|
||
def _log(msg, *args, **kwargs): # noqa: ARG001
|
||
print(f'[LOG] {msg}')
|
||
|
||
logger = MagicMock()
|
||
for method in ('info', 'debug', 'error', 'warning', 'critical',
|
||
'custom_info', 'custom_debug', 'custom_error',
|
||
'custom_warning', 'custom_critical'):
|
||
setattr(logger, method, MagicMock(side_effect=_log))
|
||
logger.base_logger = MagicMock()
|
||
return logger
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
def mock_metrics_controller(mock_logger):
|
||
"""Real MetricsController backed by the mock logger."""
|
||
return MetricsController(logger=mock_logger)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Real application fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.fixture
|
||
def notification_handler(mongodb_container, mock_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=mock_logger,
|
||
project_name='model-manager-e2e',
|
||
)
|
||
yield handler
|
||
handler.shutdown()
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
def plugin_store(gitea_container, mock_logger, mock_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=mock_logger,
|
||
notification_handler=notification_handler,
|
||
metrics_controller=mock_metrics_controller,
|
||
)
|
||
yield store
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
def test_activities(
|
||
postgres_container,
|
||
minio_container,
|
||
mlflow_tracking_dir, # noqa: ARG001 – ensures MLflow URI is set
|
||
plugin_store,
|
||
mock_logger,
|
||
notification_handler,
|
||
mock_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=mock_logger,
|
||
notification_handler=notification_handler,
|
||
metrics_controller=mock_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():
|
||
"""In-memory Temporal 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 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
|