ba9eb3d7c7
- Made `date_column` a required field in `TrainModelParams`, ensuring it must be present in the input data. - Updated related documentation in `input-sample.md`, `README.md`, and various test scenarios to reflect the change in requirement. - Adjusted the handling of `date_format` to default to `yyyy-MM-dd HH:mm:ss` if omitted, enhancing usability. - Refined test scenarios to include new examples and ensure compliance with the updated parameter structure. These changes improve the robustness of the model training workflow and clarify the expectations for input data.
650 lines
21 KiB
Python
650 lines
21 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 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
|
||
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.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
|
||
"""
|
||
|
||
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 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_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',
|
||
)
|
||
|
||
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_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'))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Observability (real sientia_do implementations)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.fixture(scope='session')
|
||
def e2e_logger():
|
||
"""Shared production-style Logger for the whole E2E session."""
|
||
return get_logger('model-manager-e2e')
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
def metrics_controller(e2e_logger):
|
||
"""MetricsController bound to the E2E logger (fresh instance per test)."""
|
||
return MetricsController(logger=e2e_logger)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Application fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest_asyncio.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_asyncio.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_asyncio.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
|