feat: enhance E2E testing setup and model reporting

- 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.
This commit is contained in:
vitor-aignosi
2026-05-05 10:59:51 -03:00
parent ba9eb3d7c7
commit d1f9394879
29 changed files with 255 additions and 182 deletions

View File

@@ -54,6 +54,8 @@ These scenarios test the business rule validations inside `validate_train_params
| `test_scenario_2_1_5_unknown_model_type` | `model_type = 'totally_unknown_model'` | `'totally_unknown_model'` | | `test_scenario_2_1_5_unknown_model_type` | `model_type = 'totally_unknown_model'` | `'totally_unknown_model'` |
| `test_scenario_2_1_6_missing_target_variable` | `target_variable = ''` | `'target_variable'` | | `test_scenario_2_1_6_missing_target_variable` | `target_variable = ''` | `'target_variable'` |
| `test_scenario_2_1_7_missing_experiment_run_id` | Missing `experiment_run_id` | N/A (raises ValueError immediately) | | `test_scenario_2_1_7_missing_experiment_run_id` | Missing `experiment_run_id` | N/A (raises ValueError immediately) |
| `test_scenario_2_1_8_missing_date_column` | Missing `date_column` | N/A (raises ValueError immediately) |
| `test_scenario_2_1_9_whitespace_date_column` | `date_column = ' '` | `'date_column'` |
## 3. CleanupFiles Workflow (`test_cleanup_files_workflow.py`) ## 3. CleanupFiles Workflow (`test_cleanup_files_workflow.py`)

View File

@@ -17,6 +17,7 @@ from concurrent.futures import ThreadPoolExecutor
import base64 import base64
import csv import csv
import io import io
import os
import shutil import shutil
import tempfile import tempfile
import time import time
@@ -27,10 +28,10 @@ import pytest_asyncio
import requests import requests
from minio import Minio from minio import Minio
from sqlalchemy import create_engine, text from sqlalchemy import create_engine, text
from testcontainers.core.container import DockerContainer from testcontainers.core.container import DockerContainer # type: ignore[import-untyped]
from testcontainers.minio import MinioContainer from testcontainers.minio import MinioContainer # type: ignore[import-untyped]
from testcontainers.mongodb import MongoDbContainer from testcontainers.mongodb import MongoDbContainer # type: ignore[import-untyped]
from testcontainers.postgres import PostgresContainer from testcontainers.postgres import PostgresContainer # type: ignore[import-untyped]
from temporalio.testing import WorkflowEnvironment from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker from temporalio.worker import Worker
@@ -288,6 +289,12 @@ class DummyTransformer:
pass pass
""" """
# requirements.txt required by current model packaging path in training activity.
requirements_txt = """
pandas
numpy
"""
def push_file(path: str, content: str): def push_file(path: str, content: str):
encoded = base64.b64encode(content.encode()).decode() encoded = base64.b64encode(content.encode()).decode()
_gitea_api( _gitea_api(
@@ -307,6 +314,7 @@ class DummyTransformer:
push_file(f'{prefix}/schemas.yaml', schemas_yaml) push_file(f'{prefix}/schemas.yaml', schemas_yaml)
push_file(f'{prefix}/wrapper.py', wrapper_py) push_file(f'{prefix}/wrapper.py', wrapper_py)
push_file(f'{prefix}/model_logic.py', model_logic_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_file(f'{prefix}/__init__.py', "")
# Push runtime # Push runtime
@@ -317,7 +325,7 @@ class DummyTransformer:
# Session-scoped containers # Session-scoped containers
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@pytest_asyncio.fixture(scope='session') @pytest.fixture(scope='session')
def postgres_container(): def postgres_container():
"""PostgreSQL 15 container for experiment_run table.""" """PostgreSQL 15 container for experiment_run table."""
container = PostgresContainer('postgres:15') container = PostgresContainer('postgres:15')
@@ -326,7 +334,7 @@ def postgres_container():
container.stop() container.stop()
@pytest_asyncio.fixture(scope='session') @pytest.fixture(scope='session')
def minio_container(): def minio_container():
"""MinIO container for training CSV storage.""" """MinIO container for training CSV storage."""
container = MinioContainer() container = MinioContainer()
@@ -335,7 +343,7 @@ def minio_container():
container.stop() container.stop()
@pytest_asyncio.fixture(scope='session') @pytest.fixture(scope='session')
def mongodb_container(): def mongodb_container():
"""MongoDB container for CoreNotificationHandler.""" """MongoDB container for CoreNotificationHandler."""
container = MongoDbContainer('mongo:7') container = MongoDbContainer('mongo:7')
@@ -344,7 +352,7 @@ def mongodb_container():
container.stop() container.stop()
@pytest_asyncio.fixture(scope='session') @pytest.fixture(scope='session')
def gitea_container(): def gitea_container():
""" """
Gitea container with a ``model-store`` repo seeded via REST API Gitea container with a ``model-store`` repo seeded via REST API
@@ -397,7 +405,7 @@ def gitea_container():
container.stop() container.stop()
@pytest_asyncio.fixture(scope='session') @pytest.fixture(scope='session')
def mlflow_tracking_dir(): def mlflow_tracking_dir():
"""Local MLflow filesystem tracking directory (no network needed).""" """Local MLflow filesystem tracking directory (no network needed)."""
tmpdir = tempfile.mkdtemp(prefix='mlflow-e2e-') tmpdir = tempfile.mkdtemp(prefix='mlflow-e2e-')
@@ -406,11 +414,51 @@ def mlflow_tracking_dir():
shutil.rmtree(tmpdir, ignore_errors=True) 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 # Session-scoped: seed MinIO with training CSV
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@pytest_asyncio.fixture(scope='session', autouse=True) @pytest.fixture(scope='session', autouse=True)
def upload_training_csv(minio_container, mlflow_tracking_dir): # noqa: ARG001 def upload_training_csv(minio_container, mlflow_tracking_dir): # noqa: ARG001
""" """
Upload training CSV files to the MinIO container before any test runs. Upload training CSV files to the MinIO container before any test runs.
@@ -479,7 +527,7 @@ def upload_training_csv(minio_container, mlflow_tracking_dir): # noqa: ARG001
# Function-scoped: database engine + schema setup # Function-scoped: database engine + schema setup
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@pytest_asyncio.fixture @pytest.fixture
def postgres_engine(postgres_container): def postgres_engine(postgres_container):
"""SQLAlchemy engine connected to the test PostgreSQL container.""" """SQLAlchemy engine connected to the test PostgreSQL container."""
engine = create_engine(postgres_container.get_connection_url()) engine = create_engine(postgres_container.get_connection_url())
@@ -487,7 +535,7 @@ def postgres_engine(postgres_container):
engine.dispose() engine.dispose()
@pytest_asyncio.fixture(autouse=True) @pytest.fixture(autouse=True)
def setup_experiment_run_table(postgres_engine): def setup_experiment_run_table(postgres_engine):
""" """
Create the experiment_run table before each test and drop it afterwards Create the experiment_run table before each test and drop it afterwards
@@ -517,13 +565,13 @@ def setup_experiment_run_table(postgres_engine):
# Observability (real sientia_do implementations) # Observability (real sientia_do implementations)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@pytest_asyncio.fixture(scope='session') @pytest.fixture(scope='session')
def e2e_logger(): def e2e_logger():
"""Shared production-style Logger for the whole E2E session.""" """Shared production-style Logger for the whole E2E session."""
return get_logger('model-manager-e2e') return get_logger('model-manager-e2e')
@pytest_asyncio.fixture @pytest.fixture
def metrics_controller(e2e_logger): def metrics_controller(e2e_logger):
"""MetricsController bound to the E2E logger (fresh instance per test).""" """MetricsController bound to the E2E logger (fresh instance per test)."""
return MetricsController(logger=e2e_logger) return MetricsController(logger=e2e_logger)
@@ -533,7 +581,7 @@ def metrics_controller(e2e_logger):
# Application fixtures # Application fixtures
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@pytest_asyncio.fixture @pytest.fixture
def notification_handler(mongodb_container, e2e_logger): def notification_handler(mongodb_container, e2e_logger):
""" """
Real CoreNotificationHandler connected to the MongoDB testcontainer. Real CoreNotificationHandler connected to the MongoDB testcontainer.
@@ -549,7 +597,7 @@ def notification_handler(mongodb_container, e2e_logger):
handler.shutdown() handler.shutdown()
@pytest_asyncio.fixture @pytest.fixture
def plugin_store(gitea_container, e2e_logger, metrics_controller, notification_handler): def plugin_store(gitea_container, e2e_logger, metrics_controller, notification_handler):
""" """
Real PluginStore pointed at the Gitea testcontainer. Real PluginStore pointed at the Gitea testcontainer.
@@ -569,7 +617,7 @@ def plugin_store(gitea_container, e2e_logger, metrics_controller, notification_h
yield store yield store
@pytest_asyncio.fixture @pytest.fixture
def test_activities( def test_activities(
postgres_container, postgres_container,
minio_container, minio_container,

View File

@@ -246,3 +246,53 @@ async def test_scenario_2_1_7_missing_experiment_run_id(
) )
combined = _exception_chain_text(excinfo.value) combined = _exception_chain_text(excinfo.value)
assert 'experiment_run_id' in combined assert 'experiment_run_id' in combined
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_8_missing_date_column(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
):
"""Scenario 2.1.8 date_column missing in payload raises before workflow business validation."""
scenario = load_scenario('01-linear-regression-basic.json')
scenario = {k: v for k, v in scenario.items() if k != 'date_column'}
with pytest.raises(WorkflowFailureError) as excinfo:
await start_and_await_workflow(
temporal_test_env.client,
TrainModel.run,
scenario,
make_workflow_id('test-s2-1-8'),
)
combined = _exception_chain_text(excinfo.value)
assert 'date_column' in combined
@pytest.mark.asyncio
@pytest.mark.integration
async def test_scenario_2_1_9_whitespace_date_column(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
postgres_engine,
):
"""Scenario 2.1.9 date_column=' ' must produce ORCHESTRATOR_VALIDATION_ERROR."""
experiment_run_id = _VALIDATION_ID_BASE + 9
scenario = load_scenario('01-linear-regression-basic.json')
scenario = {**scenario, 'experiment_run_id': experiment_run_id, 'date_column': ' '}
insert_experiment_run(postgres_engine, experiment_run_id)
with pytest.raises(Exception):
await start_and_await_workflow(
temporal_test_env.client,
TrainModel.run,
scenario,
make_workflow_id('test-s2-1-9'),
)
assert_experiment_error(
postgres_engine,
experiment_run_id,
expected_status='ORCHESTRATOR_VALIDATION_ERROR',
error_substr='date_column',
)

View File

@@ -111,8 +111,14 @@ class Reports:
] ]
self.metrics.extend(metrics) self.metrics.extend(metrics)
if run: if run:
mapping = ColumnMapping()
mapping.target = self.target_name
report = Report(metrics=metrics, options=self.options) report = Report(metrics=metrics, options=self.options)
report.run(reference_data=self.ref_data, current_data=self.cur_data) report.run(
reference_data=self.ref_data,
current_data=self.cur_data,
column_mapping=mapping,
)
self.sections['data_quality'] = report.as_dict() self.sections['data_quality'] = report.as_dict()
if self.base_path: if self.base_path:
# Note: Relies on Evidently's save_html() to properly manage file I/O # Note: Relies on Evidently's save_html() to properly manage file I/O
@@ -128,8 +134,14 @@ class Reports:
""" """
self.metrics.append(DataDriftPreset(columns=columns)) self.metrics.append(DataDriftPreset(columns=columns))
if run: if run:
mapping = ColumnMapping()
mapping.target = self.target_name
report = Report(metrics=[DataDriftPreset(columns=columns)], options=self.options) report = Report(metrics=[DataDriftPreset(columns=columns)], options=self.options)
report.run(reference_data=self.ref_data, current_data=self.cur_data) report.run(
reference_data=self.ref_data,
current_data=self.cur_data,
column_mapping=mapping,
)
self.sections['data_drift'] = report.as_dict() self.sections['data_drift'] = report.as_dict()
if self.base_path: if self.base_path:
# Note: Relies on Evidently's save_html() to properly manage file I/O # Note: Relies on Evidently's save_html() to properly manage file I/O

View File

@@ -565,6 +565,13 @@ class DataManagerRepository(SientiaMonitoring):
current_data = y_val_pred[['prediction']].join(data.val_data, how='inner') current_data = y_val_pred[['prediction']].join(data.val_data, how='inner')
current_data_float = current_data.astype(np.float64) current_data_float = current_data.astype(np.float64)
# Evidently's ConflictTargetMetric expects a literal `target` column name.
# Keep the original target column and provide this alias for report metrics.
target_col = data.params.target_variable
reference_data_float['target'] = reference_data_float[target_col]
current_data_float['target'] = current_data_float[target_col]
# Initialize report generator # Initialize report generator
base_path = self._get_reports_directory() base_path = self._get_reports_directory()
data.run_dir = self._create_run_directory(base_path, data.run_name) data.run_dir = self._create_run_directory(base_path, data.run_name)
@@ -581,7 +588,6 @@ class DataManagerRepository(SientiaMonitoring):
) )
# Generate report sections # Generate report sections
target_col = data.params.target_variable
feature_and_target_cols = data.params.variable_columns + [target_col] feature_and_target_cols = data.params.variable_columns + [target_col]
report.add_data_quality_section(columns=feature_and_target_cols) report.add_data_quality_section(columns=feature_and_target_cols)
report.add_data_drift_section(columns=feature_and_target_cols) report.add_data_drift_section(columns=feature_and_target_cols)

View File

@@ -17,6 +17,7 @@ with workflow.unsafe.imports_passed_through():
from model_manager.workflows.train_model import no_retry_policy from model_manager.workflows.train_model import no_retry_policy
TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120')) TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120'))
POD_ID = os.getenv('POD_ID')
@workflow.defn(name='cleanup_files') @workflow.defn(name='cleanup_files')
@@ -46,7 +47,7 @@ class CleanupFiles:
# Metadata for tracking # Metadata for tracking
metadata = { metadata = {
'metadata': { 'metadata': {
'pod_id': os.getenv('POD_ID'), 'pod_id': POD_ID,
'workflow_name': 'cleanup_files', 'workflow_name': 'cleanup_files',
} }
} }

View File

@@ -17,6 +17,7 @@ with workflow.unsafe.imports_passed_through():
from typing import Any from typing import Any
from temporalio.common import RetryPolicy from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError
from model_manager.activities.activities import Activities from model_manager.activities.activities import Activities
from model_manager.activities.experiment_tracking import UpdateType from model_manager.activities.experiment_tracking import UpdateType
@@ -95,7 +96,11 @@ class TrainModel:
""" """
workflow.logger.info(f'Starting train_model workflow for {input_data}') workflow.logger.info(f'Starting train_model workflow for {input_data}')
try:
experiment_run_id = self._validate_experiment_run_id(input_data) experiment_run_id = self._validate_experiment_run_id(input_data)
except ValueError as exc:
# Prevent workflow-task retries on deterministic input contract violations.
raise ApplicationError(str(exc), non_retryable=True) from exc
input_data = {**input_data, 'experiment_run_id': experiment_run_id} input_data = {**input_data, 'experiment_run_id': experiment_run_id}
model_name = input_data.get('model_name') model_name = input_data.get('model_name')

View File

View File

@@ -1 +0,0 @@
1777925550

View File

@@ -1,14 +0,0 @@
name: "linear_regression"
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

View File

@@ -1,8 +0,0 @@
class DummyModel:
def __init__(self, **kwargs):
pass
class DummyTransformer:
def __init__(self, **kwargs):
pass

View File

@@ -1,10 +0,0 @@
model:
type: object
properties: {}
data_model:
type: object
properties: {}
opt_params:
type: object
properties: {}

View File

@@ -1,27 +0,0 @@
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

View File

@@ -1 +0,0 @@
1777925551

View File

@@ -1,14 +0,0 @@
name: "polynomial_regression"
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

View File

@@ -1,8 +0,0 @@
class DummyModel:
def __init__(self, **kwargs):
pass
class DummyTransformer:
def __init__(self, **kwargs):
pass

View File

@@ -1,10 +0,0 @@
model:
type: object
properties: {}
data_model:
type: object
properties: {}
opt_params:
type: object
properties: {}

View File

@@ -1,27 +0,0 @@
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

View File

@@ -125,7 +125,7 @@ def test_cleanup_temp_directories_nonexistent_path(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.warning = MagicMock() cleanup.warning = MagicMock()
cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}}) cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}})
@@ -152,7 +152,7 @@ def test_cleanup_temp_directories_success_with_deletions(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000') old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}') old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
@@ -187,7 +187,7 @@ def test_cleanup_temp_directories_dry_run(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000') old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}') old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
@@ -217,7 +217,7 @@ def test_cleanup_temp_directories_delete_error(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.error = MagicMock() cleanup.error = MagicMock()
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000') old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
@@ -276,7 +276,7 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.debug = MagicMock() cleanup.debug = MagicMock()
# Create a file and a directory with a non-matching name # Create a file and a directory with a non-matching name
@@ -310,7 +310,7 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.error = MagicMock() cleanup.error = MagicMock()
# Create a directory with a malformed timestamp that matches the regex but fails parsing # Create a directory with a malformed timestamp that matches the regex but fails parsing
@@ -340,7 +340,7 @@ def test_cleanup_temp_directories_generic_exception(
notification_handler=mock_notification_handler, notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller, metrics_controller=mock_metrics_controller,
) )
cleanup._emit_metrics = MagicMock() cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.send_notification = MagicMock() cleanup.send_notification = MagicMock()
with patch('os.listdir', side_effect=Exception('Unexpected OS Error')): with patch('os.listdir', side_effect=Exception('Unexpected OS Error')):

View File

@@ -136,6 +136,11 @@ def test_experiment_tracking_del_with_engine_exception(
et.engine = MagicMock() et.engine = MagicMock()
class MockSuperWithError: class MockSuperWithError:
_should_raise: bool
def __init__(self) -> None:
self._should_raise = False
def __del__(self): def __del__(self):
# Only raise error if not being cleaned up by garbage collector # Only raise error if not being cleaned up by garbage collector
# This prevents the PytestUnraisableExceptionWarning # This prevents the PytestUnraisableExceptionWarning
@@ -215,7 +220,7 @@ def test_update_experiment_run_status_success(
mock_execute(*args, **kwargs) mock_execute(*args, **kwargs)
return {'rowcount': 1} return {'rowcount': 1}
et._execute_update = mock_execute_update et._execute_update = mock_execute_update # type: ignore[method-assign]
et.info = MagicMock() et.info = MagicMock()
input_data = { input_data = {
@@ -293,7 +298,7 @@ def test_update_experiment_run_status_with_error_success(
mock_execute(*args, **kwargs) mock_execute(*args, **kwargs)
return {'rowcount': 1} return {'rowcount': 1}
et._execute_update = mock_execute_update et._execute_update = mock_execute_update # type: ignore[method-assign]
et.info = MagicMock() et.info = MagicMock()
input_data = { input_data = {
@@ -339,7 +344,7 @@ def test_update_experiment_run_status_with_error_truncate_message(
mock_execute(*args, **kwargs) mock_execute(*args, **kwargs)
return {'rowcount': 1} return {'rowcount': 1}
et._execute_update = mock_execute_update et._execute_update = mock_execute_update # type: ignore[method-assign]
et.info = MagicMock() et.info = MagicMock()
long_error = 'x' * 2000 long_error = 'x' * 2000
@@ -416,7 +421,7 @@ def test_update_experiment_run_model_saved_success(
mock_execute(*args, **kwargs) mock_execute(*args, **kwargs)
return {'rowcount': 1} return {'rowcount': 1}
et._execute_update = mock_execute_update et._execute_update = mock_execute_update # type: ignore[method-assign]
et.info = MagicMock() et.info = MagicMock()
input_data = { input_data = {
@@ -526,7 +531,7 @@ def test_update_experiment_run_no_rows_updated(
def mock_execute_update(*args, **kwargs): def mock_execute_update(*args, **kwargs):
return {'rowcount': 0} return {'rowcount': 0}
et._execute_update = mock_execute_update et._execute_update = mock_execute_update # type: ignore[method-assign]
et.send_notification = MagicMock() et.send_notification = MagicMock()
input_data = { input_data = {

View File

@@ -131,13 +131,15 @@ def test_train_model_success_serializes_result(mock_mlflow, training):
training.minio_repository.download_file = MagicMock(return_value=b'csv') training.minio_repository.download_file = MagicMock(return_value=b'csv')
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
def _set_metrics(x, _w, **_kw):
x.mse_val = 0.1
x.mae_val = 0.2
x.r2_val = 0.9
return x
training.data_manager_repository.compute_regression_metrics = MagicMock( training.data_manager_repository.compute_regression_metrics = MagicMock(
side_effect=lambda x, _w, **_kw: ( side_effect=_set_metrics
setattr(x, 'mse_val', 0.1),
setattr(x, 'mae_val', 0.2),
setattr(x, 'r2_val', 0.9),
x,
)[-1]
) )
def _fill_report(x, **_kw): def _fill_report(x, **_kw):

View File

@@ -21,7 +21,7 @@ def _stub_evidently() -> None:
sys.modules['evidently'] = ev sys.modules['evidently'] = ev
mp = ModuleType('evidently.metric_preset') mp = ModuleType('evidently.metric_preset')
mp.DataDriftPreset = _make_dummy('DataDriftPreset') mp.DataDriftPreset = _make_dummy('DataDriftPreset') # type: ignore[attr-defined]
sys.modules['evidently.metric_preset'] = mp sys.modules['evidently.metric_preset'] = mp
metrics = ModuleType('evidently.metrics') metrics = ModuleType('evidently.metrics')
@@ -47,22 +47,22 @@ def _stub_evidently() -> None:
def generate_column_metrics(*_a, **_k): def generate_column_metrics(*_a, **_k):
return [] return []
base.generate_column_metrics = generate_column_metrics base.generate_column_metrics = generate_column_metrics # type: ignore[attr-defined]
sys.modules['evidently.metrics.base_metric'] = base sys.modules['evidently.metrics.base_metric'] = base
opt = ModuleType('evidently.options') opt = ModuleType('evidently.options')
opt.ColorOptions = _make_dummy('ColorOptions') opt.ColorOptions = _make_dummy('ColorOptions') # type: ignore[attr-defined]
sys.modules['evidently.options'] = opt sys.modules['evidently.options'] = opt
pipeline = ModuleType('evidently.pipeline') pipeline = ModuleType('evidently.pipeline')
sys.modules['evidently.pipeline'] = pipeline sys.modules['evidently.pipeline'] = pipeline
colmap = ModuleType('evidently.pipeline.column_mapping') colmap = ModuleType('evidently.pipeline.column_mapping')
colmap.ColumnMapping = _make_dummy('ColumnMapping') colmap.ColumnMapping = _make_dummy('ColumnMapping') # type: ignore[attr-defined]
sys.modules['evidently.pipeline.column_mapping'] = colmap sys.modules['evidently.pipeline.column_mapping'] = colmap
rep = ModuleType('evidently.report') rep = ModuleType('evidently.report')
rep.Report = _make_dummy('Report') rep.Report = _make_dummy('Report') # type: ignore[attr-defined]
sys.modules['evidently.report'] = rep sys.modules['evidently.report'] = rep
@@ -82,9 +82,9 @@ def pytest_configure(config) -> None: # noqa: ARG001
def treat_nan(input_data, *_a, **_k): def treat_nan(input_data, *_a, **_k):
return input_data return input_data
df_pre.create_features = create_features df_pre.create_features = create_features # type: ignore[attr-defined]
df_pre.limit_dataset = limit_dataset df_pre.limit_dataset = limit_dataset # type: ignore[attr-defined]
df_pre.treat_nan = treat_nan df_pre.treat_nan = treat_nan # type: ignore[attr-defined]
sys.modules['sientia_do.operations.df_preprocessor'] = df_pre sys.modules['sientia_do.operations.df_preprocessor'] = df_pre
sys.modules.setdefault('sientia_do.operations', ModuleType('sientia_do.operations')) sys.modules.setdefault('sientia_do.operations', ModuleType('sientia_do.operations'))
@@ -97,7 +97,7 @@ def pytest_configure(config) -> None: # noqa: ARG001
pass pass
ts_an.TimeSeriesDiscontinuityAnalyzer = TimeSeriesDiscontinuityAnalyzer ts_an.TimeSeriesDiscontinuityAnalyzer = TimeSeriesDiscontinuityAnalyzer # type: ignore[attr-defined]
sys.modules['sientia_do.timeseries.analyzer'] = ts_an sys.modules['sientia_do.timeseries.analyzer'] = ts_an
sys.modules.setdefault('sientia_do.timeseries', ModuleType('sientia_do.timeseries')) sys.modules.setdefault('sientia_do.timeseries', ModuleType('sientia_do.timeseries'))

View File

@@ -128,7 +128,10 @@ def test_add_data_quality_section_with_run(monkeypatch, tmp_path, stub_color_opt
ReportMock.assert_called_once_with( ReportMock.assert_called_once_with(
metrics=[summary, column_metrics, conflict, correlations], options=report.options metrics=[summary, column_metrics, conflict, correlations], options=report.options
) )
report_instance.run.assert_called_once_with(reference_data='ref', current_data='cur') run_kwargs = report_instance.run.call_args.kwargs
assert run_kwargs['reference_data'] == 'ref'
assert run_kwargs['current_data'] == 'cur'
assert run_kwargs['column_mapping'].target == 'target'
report_instance.save_html.assert_called_once_with( report_instance.save_html.assert_called_once_with(
os.path.join(str(tmp_path), 'data_quality.html') os.path.join(str(tmp_path), 'data_quality.html')
) )
@@ -160,6 +163,34 @@ def test_add_data_quality_section_run_without_base_path(monkeypatch, stub_color_
report_instance.save_html.assert_not_called() report_instance.save_html.assert_not_called()
def test_add_data_quality_section_non_default_target_keeps_conflict_metric(
monkeypatch, stub_color_options
):
summary = object()
column_metrics = object()
conflict = object()
correlations = object()
monkeypatch.setattr(reports, 'DatasetSummaryMetric', lambda: summary)
monkeypatch.setattr(
reports,
'generate_column_metrics',
lambda *args, **kwargs: column_metrics,
)
monkeypatch.setattr(reports, 'ConflictTargetMetric', lambda: conflict)
monkeypatch.setattr(reports, 'DatasetCorrelationsMetric', lambda: correlations)
report = reports.Reports(reference_data='ref', current_data='cur', target_name='sales')
report.add_data_quality_section(columns=['c1'], run=False)
assert report.metrics[-4:] == [
summary,
column_metrics,
conflict,
correlations,
]
def test_add_data_drift_section_paths(monkeypatch, tmp_path, stub_color_options): def test_add_data_drift_section_paths(monkeypatch, tmp_path, stub_color_options):
drift_instances = [object(), object(), object()] drift_instances = [object(), object(), object()]
DataDriftPresetMock = MagicMock(side_effect=drift_instances) DataDriftPresetMock = MagicMock(side_effect=drift_instances)
@@ -180,7 +211,10 @@ def test_add_data_drift_section_paths(monkeypatch, tmp_path, stub_color_options)
report.add_data_drift_section(columns=['c1'], run=True) report.add_data_drift_section(columns=['c1'], run=True)
assert report.sections['data_drift'] == {'result': 'data_drift'} assert report.sections['data_drift'] == {'result': 'data_drift'}
ReportMock.assert_called_with(metrics=[drift_instances[2]], options=report.options) ReportMock.assert_called_with(metrics=[drift_instances[2]], options=report.options)
report_instance.run.assert_called_with(reference_data='ref', current_data='cur') run_kwargs = report_instance.run.call_args.kwargs
assert run_kwargs['reference_data'] == 'ref'
assert run_kwargs['current_data'] == 'cur'
assert run_kwargs['column_mapping'].target == 'target'
report_instance.save_html.assert_called_with(os.path.join(str(tmp_path), 'data_drift.html')) report_instance.save_html.assert_called_with(os.path.join(str(tmp_path), 'data_drift.html'))
@@ -273,10 +307,12 @@ def test_set_color_options_appends(monkeypatch):
report = reports.Reports(reference_data='ref', current_data='cur', target_name='target') report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
report.set_color_options(primary_color='#111', secondary_color='#222') report.set_color_options(primary_color='#111', secondary_color='#222')
assert len(report.options) == 2 options = report.options
assert options is not None
assert len(options) == 2
assert calls[0]['primary_color'] == '#0F4C81' assert calls[0]['primary_color'] == '#0F4C81'
assert calls[1]['primary_color'] == '#111' assert calls[1]['primary_color'] == '#111'
assert report.options[1]['secondary_color'] == '#222' assert options[1]['secondary_color'] == '#222'
def test_save_all_sections_html_requires_base_path(stub_color_options): def test_save_all_sections_html_requires_base_path(stub_color_options):

View File

@@ -46,7 +46,7 @@ def test_experiment_status_comparison():
assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR' assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR'
assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC' assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC'
assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS' assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS'
assert ExperimentStatus.TRAINING_ERROR != 'TRAINING_SUCCESS' assert str(ExperimentStatus.TRAINING_ERROR) != 'TRAINING_SUCCESS'
def test_experiment_status_access_by_name(): def test_experiment_status_access_by_name():

View File

@@ -476,6 +476,35 @@ def test_generate_report_success(tmp_path):
json.load(f) json.load(f)
def test_generate_report_adds_target_alias_for_reports(tmp_path):
repo = dmr.DataManagerRepository(MagicMock())
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
tmr = TrainModelResult(
params=p,
train_data=pd.DataFrame({'v1': [1.0, 2.0], 't': [1.0, 2.0]}),
val_data=pd.DataFrame({'v1': [1.0, 2.0], 't': [1.0, 2.0]}),
y_train_pred=pd.DataFrame({'t': [1.0, 2.0]}),
y_pred=pd.DataFrame({'t': [1.0, 2.0]}),
run_name='testrun',
)
with (
patch.object(repo, '_get_reports_directory', return_value=str(tmp_path)),
patch('model_manager.utils.repository.data_manager_repository.Reports') as mrep,
):
instance = mrep.return_value
instance.save_all_sections_html = Mock()
repo.generate_report(tmr, {})
kwargs = mrep.call_args.kwargs
reference_data = kwargs['reference_data']
current_data = kwargs['current_data']
assert 'target' in reference_data.columns
assert 'target' in current_data.columns
assert reference_data['target'].equals(reference_data['t'])
assert current_data['target'].equals(current_data['t'])
def test_generate_report_skips_equation_file_when_not_linear(tmp_path): def test_generate_report_skips_equation_file_when_not_linear(tmp_path):
repo = dmr.DataManagerRepository(MagicMock()) repo = dmr.DataManagerRepository(MagicMock())
p = TrainModelParams.from_dict(_minimal_dict_for_prepare()) p = TrainModelParams.from_dict(_minimal_dict_for_prepare())

View File

@@ -1,6 +1,5 @@
"""Unit tests for the CleanupFiles workflow.""" """Unit tests for the CleanupFiles workflow."""
import os
from unittest.mock import AsyncMock, patch from unittest.mock import AsyncMock, patch
import pytest import pytest
@@ -18,7 +17,6 @@ async def test_cleanup_files_workflow(mock_workflow_module):
# Instantiate and run the workflow # Instantiate and run the workflow
workflow_instance = CleanupFiles() workflow_instance = CleanupFiles()
with patch.dict(os.environ, {'POD_ID': 'temporal-pod'}):
await workflow_instance.run({}) await workflow_instance.run({})
# Verify that the activities were called with the correct parameters # Verify that the activities were called with the correct parameters
@@ -28,7 +26,5 @@ async def test_cleanup_files_workflow(mock_workflow_module):
# Check cleanup_temp_directories call # Check cleanup_temp_directories call
local_call_args = calls[0][0][1] local_call_args = calls[0][0][1]
assert local_call_args['temp_path'] == REPORTS_TEMP_DIR assert local_call_args['temp_path'] == REPORTS_TEMP_DIR
assert local_call_args['metadata'] == { assert local_call_args['metadata']['workflow_name'] == 'cleanup_files'
'pod_id': 'temporal-pod', assert 'pod_id' in local_call_args['metadata']
'workflow_name': 'cleanup_files',
}

View File

@@ -3,6 +3,7 @@
from unittest.mock import AsyncMock, Mock, patch from unittest.mock import AsyncMock, Mock, patch
import pytest import pytest
from temporalio.exceptions import ApplicationError
from model_manager.utils.models.experiment_status import ExperimentStatus from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.models.train_model_params import TrainModelParams
@@ -295,7 +296,7 @@ async def test_run_missing_experiment_run_id(mock_wf):
from model_manager.workflows.train_model import TrainModel from model_manager.workflows.train_model import TrainModel
mock_wf.logger = Mock() mock_wf.logger = Mock()
with pytest.raises(ValueError, match='experiment_run_id is required'): with pytest.raises(ApplicationError, match='experiment_run_id is required'):
await TrainModel().run({}) await TrainModel().run({})