gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
23912d854e8e6552e1eff5366276d986569ec9ff
sientia-dataops-model-manager/e2e/conftest.py
vitor-aignosi 23912d854e chore: update dependencies and improve import statements 
- Updated `sientia_do` version to 1.12.0 in `requirements.txt` for compatibility and feature enhancements.
- Refined import statements in `conftest.py` to improve type checking, ensuring better compatibility with type checkers.
2026-05-05 15:52:39 -03:00

698 lines
23 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
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,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,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
Reference in New Issue View Git Blame Copy Permalink