From d1f93948792e4021a13b87e5599eb51b27037241 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 5 May 2026 10:59:51 -0300 Subject: [PATCH] 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. --- {e2e => docs}/scenarios.md | 2 + e2e/conftest.py | 82 +++++++++++++++---- e2e/test_train_model_validation.py | 50 +++++++++++ model_manager/sientia/reports.py | 16 +++- .../repository/data_manager_repository.py | 8 +- model_manager/workflows/cleanup_files.py | 3 +- model_manager/workflows/train_model.py | 7 +- models/__init__.py | 0 models/linear_regression/.cache_updated_at | 1 - models/linear_regression/__init__.py | 0 models/linear_regression/index.yaml | 14 ---- models/linear_regression/model_logic.py | 8 -- models/linear_regression/schemas.yaml | 10 --- models/linear_regression/wrapper.py | 27 ------ .../polynomial_regression/.cache_updated_at | 1 - models/polynomial_regression/__init__.py | 0 models/polynomial_regression/index.yaml | 14 ---- models/polynomial_regression/model_logic.py | 8 -- models/polynomial_regression/schemas.yaml | 10 --- models/polynomial_regression/wrapper.py | 27 ------ tests/activities/test_cleanup.py | 14 ++-- tests/activities/test_experiment_tracking.py | 15 ++-- tests/activities/test_training.py | 14 ++-- tests/conftest.py | 18 ++-- tests/sientia/test_reports.py | 44 +++++++++- tests/utils/models/test_experiment_status.py | 2 +- .../test_data_manager_repository.py | 29 +++++++ tests/workflows/test_cleanup_files.py | 10 +-- tests/workflows/test_train_model.py | 3 +- 29 files changed, 255 insertions(+), 182 deletions(-) rename {e2e => docs}/scenarios.md (96%) delete mode 100644 models/__init__.py delete mode 100644 models/linear_regression/.cache_updated_at delete mode 100644 models/linear_regression/__init__.py delete mode 100644 models/linear_regression/index.yaml delete mode 100644 models/linear_regression/model_logic.py delete mode 100644 models/linear_regression/schemas.yaml delete mode 100644 models/linear_regression/wrapper.py delete mode 100644 models/polynomial_regression/.cache_updated_at delete mode 100644 models/polynomial_regression/__init__.py delete mode 100644 models/polynomial_regression/index.yaml delete mode 100644 models/polynomial_regression/model_logic.py delete mode 100644 models/polynomial_regression/schemas.yaml delete mode 100644 models/polynomial_regression/wrapper.py diff --git a/e2e/scenarios.md b/docs/scenarios.md similarity index 96% rename from e2e/scenarios.md rename to docs/scenarios.md index aa0f82e..1251097 100644 --- a/e2e/scenarios.md +++ b/docs/scenarios.md @@ -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_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_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`) diff --git a/e2e/conftest.py b/e2e/conftest.py index b5da290..7f839a3 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -17,6 +17,7 @@ from concurrent.futures import ThreadPoolExecutor import base64 import csv import io +import os import shutil import tempfile import time @@ -27,10 +28,10 @@ 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 testcontainers.core.container import DockerContainer # type: ignore[import-untyped] +from testcontainers.minio import MinioContainer # type: ignore[import-untyped] +from testcontainers.mongodb import MongoDbContainer # type: ignore[import-untyped] +from testcontainers.postgres import PostgresContainer # type: ignore[import-untyped] from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker @@ -288,6 +289,12 @@ class DummyTransformer: 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( @@ -307,6 +314,7 @@ class DummyTransformer: 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 @@ -317,7 +325,7 @@ class DummyTransformer: # Session-scoped containers # --------------------------------------------------------------------------- -@pytest_asyncio.fixture(scope='session') +@pytest.fixture(scope='session') def postgres_container(): """PostgreSQL 15 container for experiment_run table.""" container = PostgresContainer('postgres:15') @@ -326,7 +334,7 @@ def postgres_container(): container.stop() -@pytest_asyncio.fixture(scope='session') +@pytest.fixture(scope='session') def minio_container(): """MinIO container for training CSV storage.""" container = MinioContainer() @@ -335,7 +343,7 @@ def minio_container(): container.stop() -@pytest_asyncio.fixture(scope='session') +@pytest.fixture(scope='session') def mongodb_container(): """MongoDB container for CoreNotificationHandler.""" container = MongoDbContainer('mongo:7') @@ -344,7 +352,7 @@ def mongodb_container(): container.stop() -@pytest_asyncio.fixture(scope='session') +@pytest.fixture(scope='session') def gitea_container(): """ Gitea container with a ``model-store`` repo seeded via REST API @@ -397,7 +405,7 @@ def gitea_container(): container.stop() -@pytest_asyncio.fixture(scope='session') +@pytest.fixture(scope='session') def mlflow_tracking_dir(): """Local MLflow filesystem tracking directory (no network needed).""" tmpdir = tempfile.mkdtemp(prefix='mlflow-e2e-') @@ -406,11 +414,51 @@ def mlflow_tracking_dir(): 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_asyncio.fixture(scope='session', autouse=True) +@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. @@ -479,7 +527,7 @@ def upload_training_csv(minio_container, mlflow_tracking_dir): # noqa: ARG001 # Function-scoped: database engine + schema setup # --------------------------------------------------------------------------- -@pytest_asyncio.fixture +@pytest.fixture def postgres_engine(postgres_container): """SQLAlchemy engine connected to the test PostgreSQL container.""" engine = create_engine(postgres_container.get_connection_url()) @@ -487,7 +535,7 @@ def postgres_engine(postgres_container): engine.dispose() -@pytest_asyncio.fixture(autouse=True) +@pytest.fixture(autouse=True) def setup_experiment_run_table(postgres_engine): """ 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) # --------------------------------------------------------------------------- -@pytest_asyncio.fixture(scope='session') +@pytest.fixture(scope='session') def e2e_logger(): """Shared production-style Logger for the whole E2E session.""" return get_logger('model-manager-e2e') -@pytest_asyncio.fixture +@pytest.fixture def metrics_controller(e2e_logger): """MetricsController bound to the E2E logger (fresh instance per test).""" return MetricsController(logger=e2e_logger) @@ -533,7 +581,7 @@ def metrics_controller(e2e_logger): # Application fixtures # --------------------------------------------------------------------------- -@pytest_asyncio.fixture +@pytest.fixture def notification_handler(mongodb_container, e2e_logger): """ Real CoreNotificationHandler connected to the MongoDB testcontainer. @@ -549,7 +597,7 @@ def notification_handler(mongodb_container, e2e_logger): handler.shutdown() -@pytest_asyncio.fixture +@pytest.fixture def plugin_store(gitea_container, e2e_logger, metrics_controller, notification_handler): """ Real PluginStore pointed at the Gitea testcontainer. @@ -569,7 +617,7 @@ def plugin_store(gitea_container, e2e_logger, metrics_controller, notification_h yield store -@pytest_asyncio.fixture +@pytest.fixture def test_activities( postgres_container, minio_container, diff --git a/e2e/test_train_model_validation.py b/e2e/test_train_model_validation.py index 7e58c3c..20e0031 100644 --- a/e2e/test_train_model_validation.py +++ b/e2e/test_train_model_validation.py @@ -246,3 +246,53 @@ async def test_scenario_2_1_7_missing_experiment_run_id( ) combined = _exception_chain_text(excinfo.value) 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', + ) diff --git a/model_manager/sientia/reports.py b/model_manager/sientia/reports.py index a425f6e..591c77d 100644 --- a/model_manager/sientia/reports.py +++ b/model_manager/sientia/reports.py @@ -111,8 +111,14 @@ class Reports: ] self.metrics.extend(metrics) if run: + mapping = ColumnMapping() + mapping.target = self.target_name 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() if self.base_path: # 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)) if run: + mapping = ColumnMapping() + mapping.target = self.target_name 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() if self.base_path: # Note: Relies on Evidently's save_html() to properly manage file I/O diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py index f4e9c70..2bcbc92 100644 --- a/model_manager/utils/repository/data_manager_repository.py +++ b/model_manager/utils/repository/data_manager_repository.py @@ -565,6 +565,13 @@ class DataManagerRepository(SientiaMonitoring): current_data = y_val_pred[['prediction']].join(data.val_data, how='inner') 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 base_path = self._get_reports_directory() data.run_dir = self._create_run_directory(base_path, data.run_name) @@ -581,7 +588,6 @@ class DataManagerRepository(SientiaMonitoring): ) # Generate report sections - target_col = data.params.target_variable feature_and_target_cols = data.params.variable_columns + [target_col] report.add_data_quality_section(columns=feature_and_target_cols) report.add_data_drift_section(columns=feature_and_target_cols) diff --git a/model_manager/workflows/cleanup_files.py b/model_manager/workflows/cleanup_files.py index d077cf5..3683f3c 100644 --- a/model_manager/workflows/cleanup_files.py +++ b/model_manager/workflows/cleanup_files.py @@ -17,6 +17,7 @@ with workflow.unsafe.imports_passed_through(): from model_manager.workflows.train_model import no_retry_policy TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120')) + POD_ID = os.getenv('POD_ID') @workflow.defn(name='cleanup_files') @@ -46,7 +47,7 @@ class CleanupFiles: # Metadata for tracking metadata = { 'metadata': { - 'pod_id': os.getenv('POD_ID'), + 'pod_id': POD_ID, 'workflow_name': 'cleanup_files', } } diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index 4e6c676..00557e7 100644 --- a/model_manager/workflows/train_model.py +++ b/model_manager/workflows/train_model.py @@ -17,6 +17,7 @@ with workflow.unsafe.imports_passed_through(): from typing import Any from temporalio.common import RetryPolicy + from temporalio.exceptions import ApplicationError from model_manager.activities.activities import Activities 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}') - experiment_run_id = self._validate_experiment_run_id(input_data) + try: + 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} model_name = input_data.get('model_name') diff --git a/models/__init__.py b/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/models/linear_regression/.cache_updated_at b/models/linear_regression/.cache_updated_at deleted file mode 100644 index c9a3a7d..0000000 --- a/models/linear_regression/.cache_updated_at +++ /dev/null @@ -1 +0,0 @@ -1777925550 \ No newline at end of file diff --git a/models/linear_regression/__init__.py b/models/linear_regression/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/models/linear_regression/index.yaml b/models/linear_regression/index.yaml deleted file mode 100644 index ae1949d..0000000 --- a/models/linear_regression/index.yaml +++ /dev/null @@ -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 diff --git a/models/linear_regression/model_logic.py b/models/linear_regression/model_logic.py deleted file mode 100644 index 1372172..0000000 --- a/models/linear_regression/model_logic.py +++ /dev/null @@ -1,8 +0,0 @@ - -class DummyModel: - def __init__(self, **kwargs): - pass - -class DummyTransformer: - def __init__(self, **kwargs): - pass diff --git a/models/linear_regression/schemas.yaml b/models/linear_regression/schemas.yaml deleted file mode 100644 index faadeed..0000000 --- a/models/linear_regression/schemas.yaml +++ /dev/null @@ -1,10 +0,0 @@ - -model: - type: object - properties: {} -data_model: - type: object - properties: {} -opt_params: - type: object - properties: {} diff --git a/models/linear_regression/wrapper.py b/models/linear_regression/wrapper.py deleted file mode 100644 index 70b32ee..0000000 --- a/models/linear_regression/wrapper.py +++ /dev/null @@ -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 diff --git a/models/polynomial_regression/.cache_updated_at b/models/polynomial_regression/.cache_updated_at deleted file mode 100644 index ffdff25..0000000 --- a/models/polynomial_regression/.cache_updated_at +++ /dev/null @@ -1 +0,0 @@ -1777925551 \ No newline at end of file diff --git a/models/polynomial_regression/__init__.py b/models/polynomial_regression/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/models/polynomial_regression/index.yaml b/models/polynomial_regression/index.yaml deleted file mode 100644 index e91d59d..0000000 --- a/models/polynomial_regression/index.yaml +++ /dev/null @@ -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 diff --git a/models/polynomial_regression/model_logic.py b/models/polynomial_regression/model_logic.py deleted file mode 100644 index 1372172..0000000 --- a/models/polynomial_regression/model_logic.py +++ /dev/null @@ -1,8 +0,0 @@ - -class DummyModel: - def __init__(self, **kwargs): - pass - -class DummyTransformer: - def __init__(self, **kwargs): - pass diff --git a/models/polynomial_regression/schemas.yaml b/models/polynomial_regression/schemas.yaml deleted file mode 100644 index faadeed..0000000 --- a/models/polynomial_regression/schemas.yaml +++ /dev/null @@ -1,10 +0,0 @@ - -model: - type: object - properties: {} -data_model: - type: object - properties: {} -opt_params: - type: object - properties: {} diff --git a/models/polynomial_regression/wrapper.py b/models/polynomial_regression/wrapper.py deleted file mode 100644 index 70b32ee..0000000 --- a/models/polynomial_regression/wrapper.py +++ /dev/null @@ -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 diff --git a/tests/activities/test_cleanup.py b/tests/activities/test_cleanup.py index 8146280..3921117 100644 --- a/tests/activities/test_cleanup.py +++ b/tests/activities/test_cleanup.py @@ -125,7 +125,7 @@ def test_cleanup_temp_directories_nonexistent_path( notification_handler=mock_notification_handler, metrics_controller=mock_metrics_controller, ) - cleanup._emit_metrics = MagicMock() + cleanup._emit_metrics = MagicMock() # type: ignore[method-assign] cleanup.warning = MagicMock() 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, 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_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, 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_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, metrics_controller=mock_metrics_controller, ) - cleanup._emit_metrics = MagicMock() + cleanup._emit_metrics = MagicMock() # type: ignore[method-assign] cleanup.error = MagicMock() 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, metrics_controller=mock_metrics_controller, ) - cleanup._emit_metrics = MagicMock() + cleanup._emit_metrics = MagicMock() # type: ignore[method-assign] cleanup.debug = MagicMock() # 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, metrics_controller=mock_metrics_controller, ) - cleanup._emit_metrics = MagicMock() + cleanup._emit_metrics = MagicMock() # type: ignore[method-assign] cleanup.error = MagicMock() # 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, metrics_controller=mock_metrics_controller, ) - cleanup._emit_metrics = MagicMock() + cleanup._emit_metrics = MagicMock() # type: ignore[method-assign] cleanup.send_notification = MagicMock() with patch('os.listdir', side_effect=Exception('Unexpected OS Error')): diff --git a/tests/activities/test_experiment_tracking.py b/tests/activities/test_experiment_tracking.py index efedd51..2851560 100644 --- a/tests/activities/test_experiment_tracking.py +++ b/tests/activities/test_experiment_tracking.py @@ -136,6 +136,11 @@ def test_experiment_tracking_del_with_engine_exception( et.engine = MagicMock() class MockSuperWithError: + _should_raise: bool + + def __init__(self) -> None: + self._should_raise = False + def __del__(self): # Only raise error if not being cleaned up by garbage collector # This prevents the PytestUnraisableExceptionWarning @@ -215,7 +220,7 @@ def test_update_experiment_run_status_success( mock_execute(*args, **kwargs) return {'rowcount': 1} - et._execute_update = mock_execute_update + et._execute_update = mock_execute_update # type: ignore[method-assign] et.info = MagicMock() input_data = { @@ -293,7 +298,7 @@ def test_update_experiment_run_status_with_error_success( mock_execute(*args, **kwargs) return {'rowcount': 1} - et._execute_update = mock_execute_update + et._execute_update = mock_execute_update # type: ignore[method-assign] et.info = MagicMock() input_data = { @@ -339,7 +344,7 @@ def test_update_experiment_run_status_with_error_truncate_message( mock_execute(*args, **kwargs) return {'rowcount': 1} - et._execute_update = mock_execute_update + et._execute_update = mock_execute_update # type: ignore[method-assign] et.info = MagicMock() long_error = 'x' * 2000 @@ -416,7 +421,7 @@ def test_update_experiment_run_model_saved_success( mock_execute(*args, **kwargs) return {'rowcount': 1} - et._execute_update = mock_execute_update + et._execute_update = mock_execute_update # type: ignore[method-assign] et.info = MagicMock() input_data = { @@ -526,7 +531,7 @@ def test_update_experiment_run_no_rows_updated( def mock_execute_update(*args, **kwargs): return {'rowcount': 0} - et._execute_update = mock_execute_update + et._execute_update = mock_execute_update # type: ignore[method-assign] et.send_notification = MagicMock() input_data = { diff --git a/tests/activities/test_training.py b/tests/activities/test_training.py index c45512c..3cd2bfc 100644 --- a/tests/activities/test_training.py +++ b/tests/activities/test_training.py @@ -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.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( - side_effect=lambda x, _w, **_kw: ( - setattr(x, 'mse_val', 0.1), - setattr(x, 'mae_val', 0.2), - setattr(x, 'r2_val', 0.9), - x, - )[-1] + side_effect=_set_metrics ) def _fill_report(x, **_kw): diff --git a/tests/conftest.py b/tests/conftest.py index f4507c7..cdeaec1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,7 +21,7 @@ def _stub_evidently() -> None: sys.modules['evidently'] = ev 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 metrics = ModuleType('evidently.metrics') @@ -47,22 +47,22 @@ def _stub_evidently() -> None: def generate_column_metrics(*_a, **_k): 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 opt = ModuleType('evidently.options') - opt.ColorOptions = _make_dummy('ColorOptions') + opt.ColorOptions = _make_dummy('ColorOptions') # type: ignore[attr-defined] sys.modules['evidently.options'] = opt pipeline = ModuleType('evidently.pipeline') sys.modules['evidently.pipeline'] = pipeline 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 rep = ModuleType('evidently.report') - rep.Report = _make_dummy('Report') + rep.Report = _make_dummy('Report') # type: ignore[attr-defined] sys.modules['evidently.report'] = rep @@ -82,9 +82,9 @@ def pytest_configure(config) -> None: # noqa: ARG001 def treat_nan(input_data, *_a, **_k): return input_data - df_pre.create_features = create_features - df_pre.limit_dataset = limit_dataset - df_pre.treat_nan = treat_nan + df_pre.create_features = create_features # type: ignore[attr-defined] + df_pre.limit_dataset = limit_dataset # type: ignore[attr-defined] + df_pre.treat_nan = treat_nan # type: ignore[attr-defined] sys.modules['sientia_do.operations.df_preprocessor'] = df_pre sys.modules.setdefault('sientia_do.operations', ModuleType('sientia_do.operations')) @@ -97,7 +97,7 @@ def pytest_configure(config) -> None: # noqa: ARG001 pass - ts_an.TimeSeriesDiscontinuityAnalyzer = TimeSeriesDiscontinuityAnalyzer + ts_an.TimeSeriesDiscontinuityAnalyzer = TimeSeriesDiscontinuityAnalyzer # type: ignore[attr-defined] sys.modules['sientia_do.timeseries.analyzer'] = ts_an sys.modules.setdefault('sientia_do.timeseries', ModuleType('sientia_do.timeseries')) diff --git a/tests/sientia/test_reports.py b/tests/sientia/test_reports.py index 46fcd0d..27bd264 100644 --- a/tests/sientia/test_reports.py +++ b/tests/sientia/test_reports.py @@ -128,7 +128,10 @@ def test_add_data_quality_section_with_run(monkeypatch, tmp_path, stub_color_opt ReportMock.assert_called_once_with( 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( 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() +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): drift_instances = [object(), object(), object()] 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) assert report.sections['data_drift'] == {'result': 'data_drift'} 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')) @@ -273,10 +307,12 @@ def test_set_color_options_appends(monkeypatch): report = reports.Reports(reference_data='ref', current_data='cur', target_name='target') 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[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): diff --git a/tests/utils/models/test_experiment_status.py b/tests/utils/models/test_experiment_status.py index 42682d9..da4d7cc 100644 --- a/tests/utils/models/test_experiment_status.py +++ b/tests/utils/models/test_experiment_status.py @@ -46,7 +46,7 @@ def test_experiment_status_comparison(): assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR' assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC' 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(): diff --git a/tests/utils/repository/test_data_manager_repository.py b/tests/utils/repository/test_data_manager_repository.py index 67c25d1..ee1b4ce 100644 --- a/tests/utils/repository/test_data_manager_repository.py +++ b/tests/utils/repository/test_data_manager_repository.py @@ -476,6 +476,35 @@ def test_generate_report_success(tmp_path): 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): repo = dmr.DataManagerRepository(MagicMock()) p = TrainModelParams.from_dict(_minimal_dict_for_prepare()) diff --git a/tests/workflows/test_cleanup_files.py b/tests/workflows/test_cleanup_files.py index 7f822b0..7e84bcd 100644 --- a/tests/workflows/test_cleanup_files.py +++ b/tests/workflows/test_cleanup_files.py @@ -1,6 +1,5 @@ """Unit tests for the CleanupFiles workflow.""" -import os from unittest.mock import AsyncMock, patch import pytest @@ -18,8 +17,7 @@ async def test_cleanup_files_workflow(mock_workflow_module): # Instantiate and run the workflow 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 calls = mock_workflow_module.execute_activity_method.call_args_list @@ -28,7 +26,5 @@ async def test_cleanup_files_workflow(mock_workflow_module): # Check cleanup_temp_directories call local_call_args = calls[0][0][1] assert local_call_args['temp_path'] == REPORTS_TEMP_DIR - assert local_call_args['metadata'] == { - 'pod_id': 'temporal-pod', - 'workflow_name': 'cleanup_files', - } + assert local_call_args['metadata']['workflow_name'] == 'cleanup_files' + assert 'pod_id' in local_call_args['metadata'] diff --git a/tests/workflows/test_train_model.py b/tests/workflows/test_train_model.py index b8a7e1c..2f94d5a 100644 --- a/tests/workflows/test_train_model.py +++ b/tests/workflows/test_train_model.py @@ -3,6 +3,7 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +from temporalio.exceptions import ApplicationError from model_manager.utils.models.experiment_status import ExperimentStatus 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 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({})