From d8583cdae7694e314b085496c7bd274521421f06 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Mon, 6 Oct 2025 16:49:18 -0300 Subject: [PATCH] SIENTIAPDE-1249: Implement data models for Model Manager and add unit tests. This commit introduces data transfer objects (DTOs) and model classes for experiment status, training parameters, and training results, along with corresponding unit tests to ensure their correct behavior. --- model_manager/utils/models/__init__.py | 16 + .../utils/models/experiment_status.py | 32 ++ .../utils/models/train_model_params.py | 166 ++++++++++ .../utils/models/train_model_result.py | 55 ++++ tests/laborious/utils/models/__init__.py | 0 .../utils/models/test_experiment_status.py | 79 +++++ tests/laborious/utils/models/test_init.py | 51 +++ .../utils/models/test_train_model_params.py | 295 ++++++++++++++++++ .../utils/models/test_train_model_result.py | 246 +++++++++++++++ 9 files changed, 940 insertions(+) create mode 100644 model_manager/utils/models/__init__.py create mode 100644 model_manager/utils/models/experiment_status.py create mode 100644 model_manager/utils/models/train_model_params.py create mode 100644 model_manager/utils/models/train_model_result.py create mode 100644 tests/laborious/utils/models/__init__.py create mode 100644 tests/laborious/utils/models/test_experiment_status.py create mode 100644 tests/laborious/utils/models/test_init.py create mode 100644 tests/laborious/utils/models/test_train_model_params.py create mode 100644 tests/laborious/utils/models/test_train_model_result.py diff --git a/model_manager/utils/models/__init__.py b/model_manager/utils/models/__init__.py new file mode 100644 index 0000000..6a50ae3 --- /dev/null +++ b/model_manager/utils/models/__init__.py @@ -0,0 +1,16 @@ +""" +Models and DTOs for the Model Manager system. + +This module contains data transfer objects (DTOs) and model classes used +throughout the Model Manager workflows and activities. +""" + +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_result import TrainModelResult + +__all__ = [ + 'ExperimentStatus', + 'TrainModelParams', + 'TrainModelResult', +] diff --git a/model_manager/utils/models/experiment_status.py b/model_manager/utils/models/experiment_status.py new file mode 100644 index 0000000..139f2de --- /dev/null +++ b/model_manager/utils/models/experiment_status.py @@ -0,0 +1,32 @@ +from enum import Enum + + +class ExperimentStatus(str, Enum): + """ + Status values for experiment run lifecycle. + + This enum defines all possible status values that an experiment run can have + throughout its lifecycle, from initialization through training, model saving, + and cleanup. These statuses are used to track progress and identify failures + in the training pipeline. + + The status values follow the naming convention from the original Mage pipeline + to maintain compatibility with existing database records and monitoring systems. + + Attributes: + MAGE_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing. + TRAINING_SUCCESS: Training completed successfully with model and metrics calculated. + TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions. + MLFLOW_SENT: Model successfully saved to MLFlow. + MLFLOW_SEND_ERROR: Model saving to MLFlow failed due to connection or serialization errors. + FILE_DELETED: Cleanup completed successfully with all artifacts removed. + FILE_DELETE_ERROR: Cleanup failed due to file system or MinIO errors. + """ + + MAGE_WAITING_PROC = 'MAGE_WAITING_PROC' + TRAINING_SUCCESS = 'TRAINING_SUCCESS' + TRAINING_ERROR = 'TRAINING_ERROR' + MLFLOW_SENT = 'MLFLOW_SENT' + MLFLOW_SEND_ERROR = 'MLFLOW_SEND_ERROR' + FILE_DELETED = 'FILE_DELETED' + FILE_DELETE_ERROR = 'FILE_DELETE_ERROR' diff --git a/model_manager/utils/models/train_model_params.py b/model_manager/utils/models/train_model_params.py new file mode 100644 index 0000000..f71007f --- /dev/null +++ b/model_manager/utils/models/train_model_params.py @@ -0,0 +1,166 @@ +from typing import Any + + +class TrainModelParams: + """ + Parameters for machine learning model training. + + This class encapsulates all configuration parameters required for the training + pipeline, including data processing settings, model configuration, and experiment + tracking information. All parameters are validated upon initialization to ensure + data integrity and prevent runtime errors. + + Attributes: + variable_columns (List[str]): List of variable column names to use as features. + lag_train (int): Number of lags to apply during training phase. + lag_val (int): Number of lags to apply during validation phase. + target_variable (str): Name of the target variable to predict. + rem_static_win (bool): Whether to remove static windows from data. + low_lim (Dict[str, float]): Dictionary of lower limits for each variable. + upp_lim (Dict[str, float]): Dictionary of upper limits for each variable. + window (int): Window size for rolling operations. + use_scaler (bool): Whether to use a scaler for data normalization. + include_ar (bool): Whether to include autoregressive variables. + bucket_name (str): Name of the MinIO bucket containing training data. + file_name (str): Name of the file in the MinIO bucket. + line_separator (str): Line separator used in the CSV file. + decimal_separator (str): Decimal separator used in the CSV file. + train_size (int): Percentage of data to use for training (0-100). + shuffle (bool): Whether to shuffle the data during train/test split. + experiment_run_id (int): Unique identifier for the experiment run. + experiment_name (str): Name of the experiment for tracking. + experiment_description (str): Description of the experiment. + removed_intervals (list): List of time intervals to remove from the data. + """ + + def __init__( + self, + variable_columns: Any | None, + lag_train: Any | None, + lag_val: Any | None, + target_variable: Any | None, + rem_static_win: Any | None, + low_lim: Any | None, + upp_lim: Any | None, + window: Any | None, + use_scaler: Any | None, + include_ar: Any | None, + bucket_name: Any | None, + file_name: Any | None, + line_separator: Any | None, + decimal_separator: Any | None, + train_size: Any | None, + shuffle: Any | None, + experiment_run_id: Any | None, + experiment_name: Any | None, + experiment_description: Any | None, + removed_intervals: Any | None, + ): + """ + Initialize the TrainModelParams object and validate all input parameters. + + This constructor validates that all required parameters are provided and + have the correct types. It raises descriptive errors if validation fails, + helping to catch configuration issues early in the pipeline. + + Args: + variable_columns (list | None): List of variable column names. + lag_train (int | None): Number of lags to apply during training. + lag_val (int | None): Number of lags to apply during validation. + target_variable (str | None): Name of the target variable. + rem_static_win (bool | None): Whether to remove static windows. + low_lim (dict | None): Dictionary of lower limits for variables. + upp_lim (dict | None): Dictionary of upper limits for variables. + window (int | None): Window size for rolling operations. + use_scaler (bool | None): Whether to use a scaler for normalization. + include_ar (bool | None): Whether to include autoregressive variables. + bucket_name (str | None): Name of the MinIO bucket. + file_name (str | None): Name of the file in the MinIO bucket. + line_separator (str | None): Line separator used in the file. + decimal_separator (str | None): Decimal separator used in the file. + train_size (int | None): Percentage of data to use for training. + shuffle (bool | None): Whether to shuffle the data during splitting. + experiment_run_id (int | None): ID of the experiment run. + experiment_name (str | None): Name of the experiment. + experiment_description (str | None): Description of the experiment. + removed_intervals (list | None): List of intervals to remove from the data. + + Raises: + ValueError: If any required parameter is None. + TypeError: If any parameter is not of the expected type. + """ + self.variable_columns: list[str] = self._check_none( + variable_columns, list, 'variable_columns' + ) + self.lag_train: int = self._check_none(lag_train, int, 'lag_train') + self.lag_val: int = self._check_none(lag_val, int, 'lag_val') + self.target_variable: str = self._check_none(target_variable, str, 'target_variable') + self.rem_static_win: bool = self._check_none(rem_static_win, bool, 'rem_static_win') + self.low_lim: dict[str, float] = self._check_none(low_lim, dict, 'low_lim') + self.upp_lim: dict[str, float] = self._check_none(upp_lim, dict, 'upp_lim') + self.window: int = self._check_none(window, int, 'window') + self.use_scaler: bool = self._check_none(use_scaler, bool, 'use_scaler') + self.include_ar: bool = self._check_none(include_ar, bool, 'include_ar') + self.bucket_name: str = self._check_none(bucket_name, str, 'bucket_name') + self.file_name: str = self._check_none(file_name, str, 'file_name') + self.line_separator: str = self._check_none(line_separator, str, 'line_separator') + self.decimal_separator: str = self._check_none(decimal_separator, str, 'decimal_separator') + self.train_size: int = self._check_none(train_size, int, 'train_size') + self.shuffle: bool = self._check_none(shuffle, bool, 'shuffle') + self.experiment_run_id: int = self._check_none(experiment_run_id, int, 'experiment_run_id') + self.experiment_name: str = self._check_none(experiment_name, str, 'experiment_name') + self.experiment_description: str = self._check_none( + experiment_description, str, 'experiment_description' + ) + self.removed_intervals: list = self._check_type( + removed_intervals, list, 'removed_intervals' + ) + + def _check_none(self, value: Any | None, expected_type: type, field_name: str): + """ + Validate that a value is not None and check its type. + + This method ensures that required parameters are provided and have the + correct type, raising descriptive errors if validation fails. + + Args: + value (Any | None): The value to validate. + expected_type (type): The expected type of the value. + field_name (str): The name of the field being validated (for error messages). + + Returns: + Any: The validated value if it is not None and matches the expected type. + + Raises: + ValueError: If the value is None. + TypeError: If the value is not of the expected type. + """ + if value is None: + error = f'{field_name} is required and cannot be None.' + raise ValueError(error) + + return self._check_type(value, expected_type, field_name) + + def _check_type(self, value: Any | None, expected_type: type, field_name: str): + """ + Validate that a value matches the expected type. + + This method checks type compatibility and raises a descriptive error + if the value does not match the expected type. + + Args: + value (Any | None): The value to validate. + expected_type (type): The expected type of the value. + field_name (str): The name of the field being validated (for error messages). + + Returns: + Any: The validated value if it matches the expected type. + + Raises: + TypeError: If the value is not of the expected type. + """ + if value is not None and not isinstance(value, expected_type): + error = f'{field_name} must be of type {expected_type.__name__}, but got {type(value).__name__}.' + raise TypeError(error) + + return value diff --git a/model_manager/utils/models/train_model_result.py b/model_manager/utils/models/train_model_result.py new file mode 100644 index 0000000..6a02f07 --- /dev/null +++ b/model_manager/utils/models/train_model_result.py @@ -0,0 +1,55 @@ +from dataclasses import dataclass + +import pandas as pd +from sientia.linear_models import LinearRegressionModel +from sientia.preprocessing import DataPreprocessor + +from model_manager.utils.models.train_model_params import TrainModelParams + + +@dataclass +class TrainModelResult: + """ + A data container for storing the results of a machine learning training process. + + This dataclass encapsulates all outputs from the training pipeline, including + the trained model, datasets, evaluation metrics, and paths to generated artifacts. + It is used to pass results between activities in the training workflow. + + Attributes: + params (TrainModelParams): The parameters used to train the model. + process_data (DataPreprocessor): The data preprocessor object used to process the input data. + X_train (pd.DataFrame): The training dataset features. + X_test (pd.DataFrame): The testing dataset features. + y_train (pd.DataFrame): The training dataset target values. + y_test (pd.DataFrame): The testing dataset target values. + regr (LinearRegressionModel): The trained linear regression model. + scaler_dict (dict): A dictionary containing the scalers used to scale the features and target values. + y_pred (Optional[pd.Series]): The predicted target values for the testing dataset. Default is None. + mse_val (Optional[float]): The Mean Squared Error (MSE) of the predictions. Default is None. + mae_val (Optional[float]): The Mean Absolute Error (MAE) of the predictions. Default is None. + r2_val (Optional[float]): The R-squared (R²) value of the predictions. Default is None. + run_name (Optional[str]): The name of the MLFlow run. Default is None. + report_path (Optional[str]): The path to the generated HTML report file. Default is None. + train_data_path (Optional[str]): The path to the training dataset CSV file. Default is None. + test_data_path (Optional[str]): The path to the testing dataset CSV file. Default is None. + run_dir (Optional[str]): The path to the run directory containing all artifacts. Default is None. + """ + + params: TrainModelParams + process_data: DataPreprocessor + X_train: pd.DataFrame + X_test: pd.DataFrame + y_train: pd.DataFrame + y_test: pd.DataFrame + regr: LinearRegressionModel + scaler_dict: dict + y_pred: pd.Series | None = None + mse_val: float | None = None + mae_val: float | None = None + r2_val: float | None = None + run_name: str | None = None + report_path: str | None = None + train_data_path: str | None = None + test_data_path: str | None = None + run_dir: str | None = None diff --git a/tests/laborious/utils/models/__init__.py b/tests/laborious/utils/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/utils/models/test_experiment_status.py b/tests/laborious/utils/models/test_experiment_status.py new file mode 100644 index 0000000..238b06f --- /dev/null +++ b/tests/laborious/utils/models/test_experiment_status.py @@ -0,0 +1,79 @@ +"""Unit tests for ExperimentStatus enum.""" + +from model_manager.utils.models.experiment_status import ExperimentStatus + + +def test_experiment_status_values(): + """Test that all expected status values exist.""" + assert ExperimentStatus.MAGE_WAITING_PROC == 'MAGE_WAITING_PROC' + assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS' + assert ExperimentStatus.TRAINING_ERROR == 'TRAINING_ERROR' + assert ExperimentStatus.MLFLOW_SENT == 'MLFLOW_SENT' + assert ExperimentStatus.MLFLOW_SEND_ERROR == 'MLFLOW_SEND_ERROR' + assert ExperimentStatus.FILE_DELETED == 'FILE_DELETED' + assert ExperimentStatus.FILE_DELETE_ERROR == 'FILE_DELETE_ERROR' + + +def test_experiment_status_count(): + """Test that enum has exactly 7 status values.""" + assert len(ExperimentStatus) == 7 + + +def test_experiment_status_is_string(): + """Test that enum values are strings.""" + for status in ExperimentStatus: + assert isinstance(status.value, str) + assert isinstance(status, str) + + +def test_experiment_status_membership(): + """Test membership checks for status values.""" + assert 'MAGE_WAITING_PROC' in [s.value for s in ExperimentStatus] + assert 'TRAINING_SUCCESS' in [s.value for s in ExperimentStatus] + assert 'TRAINING_ERROR' in [s.value for s in ExperimentStatus] + assert 'MLFLOW_SENT' in [s.value for s in ExperimentStatus] + assert 'MLFLOW_SEND_ERROR' in [s.value for s in ExperimentStatus] + assert 'FILE_DELETED' in [s.value for s in ExperimentStatus] + assert 'FILE_DELETE_ERROR' in [s.value for s in ExperimentStatus] + + +def test_experiment_status_iteration(): + """Test that enum can be iterated.""" + statuses = list(ExperimentStatus) + assert len(statuses) == 7 + assert ExperimentStatus.MAGE_WAITING_PROC in statuses + assert ExperimentStatus.TRAINING_SUCCESS in statuses + assert ExperimentStatus.TRAINING_ERROR in statuses + assert ExperimentStatus.MLFLOW_SENT in statuses + assert ExperimentStatus.MLFLOW_SEND_ERROR in statuses + assert ExperimentStatus.FILE_DELETED in statuses + assert ExperimentStatus.FILE_DELETE_ERROR in statuses + + +def test_experiment_status_comparison(): + """Test that enum values can be compared with strings.""" + assert ExperimentStatus.MAGE_WAITING_PROC == 'MAGE_WAITING_PROC' + assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS' + assert ExperimentStatus.TRAINING_ERROR != 'TRAINING_SUCCESS' + + +def test_experiment_status_access_by_name(): + """Test accessing enum members by name.""" + assert ExperimentStatus['MAGE_WAITING_PROC'] == ExperimentStatus.MAGE_WAITING_PROC + assert ExperimentStatus['TRAINING_SUCCESS'] == ExperimentStatus.TRAINING_SUCCESS + assert ExperimentStatus['TRAINING_ERROR'] == ExperimentStatus.TRAINING_ERROR + assert ExperimentStatus['MLFLOW_SENT'] == ExperimentStatus.MLFLOW_SENT + assert ExperimentStatus['MLFLOW_SEND_ERROR'] == ExperimentStatus.MLFLOW_SEND_ERROR + assert ExperimentStatus['FILE_DELETED'] == ExperimentStatus.FILE_DELETED + assert ExperimentStatus['FILE_DELETE_ERROR'] == ExperimentStatus.FILE_DELETE_ERROR + + +def test_experiment_status_access_by_value(): + """Test accessing enum members by value.""" + assert ExperimentStatus('MAGE_WAITING_PROC') == ExperimentStatus.MAGE_WAITING_PROC + assert ExperimentStatus('TRAINING_SUCCESS') == ExperimentStatus.TRAINING_SUCCESS + assert ExperimentStatus('TRAINING_ERROR') == ExperimentStatus.TRAINING_ERROR + assert ExperimentStatus('MLFLOW_SENT') == ExperimentStatus.MLFLOW_SENT + assert ExperimentStatus('MLFLOW_SEND_ERROR') == ExperimentStatus.MLFLOW_SEND_ERROR + assert ExperimentStatus('FILE_DELETED') == ExperimentStatus.FILE_DELETED + assert ExperimentStatus('FILE_DELETE_ERROR') == ExperimentStatus.FILE_DELETE_ERROR diff --git a/tests/laborious/utils/models/test_init.py b/tests/laborious/utils/models/test_init.py new file mode 100644 index 0000000..aa7cff1 --- /dev/null +++ b/tests/laborious/utils/models/test_init.py @@ -0,0 +1,51 @@ +"""Unit tests for models __init__.py module.""" + +from model_manager.utils.models import ( + ExperimentStatus, + TrainModelParams, + TrainModelResult, +) + + +def test_experiment_status_import(): + """Test that ExperimentStatus can be imported from models package.""" + assert ExperimentStatus is not None + assert hasattr(ExperimentStatus, 'MAGE_WAITING_PROC') + assert hasattr(ExperimentStatus, 'TRAINING_SUCCESS') + + +def test_train_model_params_import(): + """Test that TrainModelParams can be imported from models package.""" + assert TrainModelParams is not None + assert callable(TrainModelParams) + + +def test_train_model_result_import(): + """Test that TrainModelResult can be imported from models package.""" + assert TrainModelResult is not None + # Dataclasses have __dataclass_fields__ + assert hasattr(TrainModelResult, '__dataclass_fields__') + + +def test_all_exports(): + """Test that __all__ contains all expected exports.""" + from model_manager.utils.models import __all__ + + assert 'ExperimentStatus' in __all__ + assert 'TrainModelParams' in __all__ + assert 'TrainModelResult' in __all__ + assert len(__all__) == 3 + + +def test_no_extra_exports(): + """Test that only expected items are exported.""" + import model_manager.utils.models as models_module + + # Get all public attributes (not starting with _) + public_attrs = [attr for attr in dir(models_module) if not attr.startswith('_')] + + # Should only have the 3 main classes + expected_public = {'ExperimentStatus', 'TrainModelParams', 'TrainModelResult'} + + # Check that our expected classes are present + assert expected_public.issubset(set(public_attrs)) diff --git a/tests/laborious/utils/models/test_train_model_params.py b/tests/laborious/utils/models/test_train_model_params.py new file mode 100644 index 0000000..20de30a --- /dev/null +++ b/tests/laborious/utils/models/test_train_model_params.py @@ -0,0 +1,295 @@ +"""Unit tests for TrainModelParams class.""" + +import pytest + +from model_manager.utils.models.train_model_params import TrainModelParams + + +@pytest.fixture +def valid_params_dict(): + """Create valid parameters dictionary for testing.""" + return { + 'variable_columns': ['var1', 'var2', 'var3'], + 'lag_train': 5, + 'lag_val': 3, + 'target_variable': 'target', + 'rem_static_win': True, + 'low_lim': {'var1': 0.0, 'var2': 0.0, 'var3': 0.0}, + 'upp_lim': {'var1': 100.0, 'var2': 100.0, 'var3': 100.0}, + 'window': 10, + 'use_scaler': True, + 'include_ar': False, + 'bucket_name': 'test-bucket', + 'file_name': 'test-file.csv', + 'line_separator': '\n', + 'decimal_separator': '.', + 'train_size': 80, + 'shuffle': True, + 'experiment_run_id': 123, + 'experiment_name': 'test-experiment', + 'experiment_description': 'Test experiment description', + 'removed_intervals': [], + } + + +def test_train_model_params_creation_with_valid_params(valid_params_dict): + """Test creating TrainModelParams with all valid parameters.""" + params = TrainModelParams(**valid_params_dict) + + assert params.variable_columns == ['var1', 'var2', 'var3'] + assert params.lag_train == 5 + assert params.lag_val == 3 + assert params.target_variable == 'target' + assert params.rem_static_win is True + assert params.low_lim == {'var1': 0.0, 'var2': 0.0, 'var3': 0.0} + assert params.upp_lim == {'var1': 100.0, 'var2': 100.0, 'var3': 100.0} + assert params.window == 10 + assert params.use_scaler is True + assert params.include_ar is False + assert params.bucket_name == 'test-bucket' + assert params.file_name == 'test-file.csv' + assert params.line_separator == '\n' + assert params.decimal_separator == '.' + assert params.train_size == 80 + assert params.shuffle is True + assert params.experiment_run_id == 123 + assert params.experiment_name == 'test-experiment' + assert params.experiment_description == 'Test experiment description' + assert params.removed_intervals == [] + + +def test_train_model_params_variable_columns_none_raises_error(valid_params_dict): + """Test that None variable_columns raises ValueError.""" + valid_params_dict['variable_columns'] = None + + with pytest.raises(ValueError, match='variable_columns is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_variable_columns_wrong_type_raises_error(valid_params_dict): + """Test that wrong type for variable_columns raises TypeError.""" + valid_params_dict['variable_columns'] = 'not a list' + + with pytest.raises(TypeError, match='variable_columns must be of type list'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_lag_train_none_raises_error(valid_params_dict): + """Test that None lag_train raises ValueError.""" + valid_params_dict['lag_train'] = None + + with pytest.raises(ValueError, match='lag_train is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_lag_train_wrong_type_raises_error(valid_params_dict): + """Test that wrong type for lag_train raises TypeError.""" + valid_params_dict['lag_train'] = '5' + + with pytest.raises(TypeError, match='lag_train must be of type int'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_target_variable_none_raises_error(valid_params_dict): + """Test that None target_variable raises ValueError.""" + valid_params_dict['target_variable'] = None + + with pytest.raises(ValueError, match='target_variable is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_target_variable_wrong_type_raises_error(valid_params_dict): + """Test that wrong type for target_variable raises TypeError.""" + valid_params_dict['target_variable'] = 123 + + with pytest.raises(TypeError, match='target_variable must be of type str'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_boolean_fields(valid_params_dict): + """Test boolean fields validation.""" + # Test rem_static_win + valid_params_dict['rem_static_win'] = None + with pytest.raises(ValueError, match='rem_static_win is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + valid_params_dict['rem_static_win'] = 'true' + with pytest.raises(TypeError, match='rem_static_win must be of type bool'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_dict_fields(valid_params_dict): + """Test dict fields validation.""" + # Test low_lim + valid_params_dict['low_lim'] = None + with pytest.raises(ValueError, match='low_lim is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + valid_params_dict['low_lim'] = {'var1': 0.0} + valid_params_dict['upp_lim'] = 'not a dict' + with pytest.raises(TypeError, match='upp_lim must be of type dict'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_bucket_name_none_raises_error(valid_params_dict): + """Test that None bucket_name raises ValueError.""" + valid_params_dict['bucket_name'] = None + + with pytest.raises(ValueError, match='bucket_name is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_file_name_none_raises_error(valid_params_dict): + """Test that None file_name raises ValueError.""" + valid_params_dict['file_name'] = None + + with pytest.raises(ValueError, match='file_name is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_experiment_run_id_none_raises_error(valid_params_dict): + """Test that None experiment_run_id raises ValueError.""" + valid_params_dict['experiment_run_id'] = None + + with pytest.raises(ValueError, match='experiment_run_id is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_experiment_name_none_raises_error(valid_params_dict): + """Test that None experiment_name raises ValueError.""" + valid_params_dict['experiment_name'] = None + + with pytest.raises(ValueError, match='experiment_name is required and cannot be None'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_removed_intervals_can_be_none(valid_params_dict): + """Test that removed_intervals can be None (uses _check_type not _check_none).""" + valid_params_dict['removed_intervals'] = None + + params = TrainModelParams(**valid_params_dict) + assert params.removed_intervals is None + + +def test_train_model_params_removed_intervals_wrong_type_raises_error(valid_params_dict): + """Test that wrong type for removed_intervals raises TypeError.""" + valid_params_dict['removed_intervals'] = 'not a list' + + with pytest.raises(TypeError, match='removed_intervals must be of type list'): + TrainModelParams(**valid_params_dict) + + +def test_train_model_params_removed_intervals_with_values(valid_params_dict): + """Test removed_intervals with actual interval values.""" + valid_params_dict['removed_intervals'] = [ + ('2023-01-01', '2023-01-10'), + ('2023-02-01', '2023-02-05'), + ] + + params = TrainModelParams(**valid_params_dict) + assert len(params.removed_intervals) == 2 + assert params.removed_intervals[0] == ('2023-01-01', '2023-01-10') + + +def test_train_model_params_all_fields_count(): + """Test that TrainModelParams has exactly 20 required fields.""" + import inspect + + sig = inspect.signature(TrainModelParams.__init__) + # Subtract 1 for 'self' + param_count = len(sig.parameters) - 1 + assert param_count == 20 + + +def test_train_model_params_with_minimal_valid_data(): + """Test creating params with minimal valid data.""" + params = TrainModelParams( + variable_columns=['x'], + lag_train=1, + lag_val=1, + target_variable='y', + rem_static_win=False, + low_lim={}, + upp_lim={}, + window=1, + use_scaler=False, + include_ar=False, + bucket_name='bucket', + file_name='file.csv', + line_separator='\n', + decimal_separator='.', + train_size=50, + shuffle=False, + experiment_run_id=1, + experiment_name='exp', + experiment_description='desc', + removed_intervals=[], + ) + + assert params.variable_columns == ['x'] + assert params.lag_train == 1 + assert params.experiment_run_id == 1 + + +def test_train_model_params_check_none_method(): + """Test _check_none method behavior.""" + params_dict = { + 'variable_columns': ['var1'], + 'lag_train': 5, + 'lag_val': 3, + 'target_variable': 'target', + 'rem_static_win': True, + 'low_lim': {}, + 'upp_lim': {}, + 'window': 10, + 'use_scaler': True, + 'include_ar': False, + 'bucket_name': 'bucket', + 'file_name': 'file.csv', + 'line_separator': '\n', + 'decimal_separator': '.', + 'train_size': 80, + 'shuffle': True, + 'experiment_run_id': 123, + 'experiment_name': 'exp', + 'experiment_description': 'desc', + 'removed_intervals': [], + } + + params = TrainModelParams(**params_dict) + + # Test that _check_none is a private method + assert hasattr(params, '_check_none') + assert callable(params._check_none) + + +def test_train_model_params_check_type_method(): + """Test _check_type method behavior.""" + params_dict = { + 'variable_columns': ['var1'], + 'lag_train': 5, + 'lag_val': 3, + 'target_variable': 'target', + 'rem_static_win': True, + 'low_lim': {}, + 'upp_lim': {}, + 'window': 10, + 'use_scaler': True, + 'include_ar': False, + 'bucket_name': 'bucket', + 'file_name': 'file.csv', + 'line_separator': '\n', + 'decimal_separator': '.', + 'train_size': 80, + 'shuffle': True, + 'experiment_run_id': 123, + 'experiment_name': 'exp', + 'experiment_description': 'desc', + 'removed_intervals': [], + } + + params = TrainModelParams(**params_dict) + + # Test that _check_type is a private method + assert hasattr(params, '_check_type') + assert callable(params._check_type) diff --git a/tests/laborious/utils/models/test_train_model_result.py b/tests/laborious/utils/models/test_train_model_result.py new file mode 100644 index 0000000..b99081b --- /dev/null +++ b/tests/laborious/utils/models/test_train_model_result.py @@ -0,0 +1,246 @@ +"""Unit tests for TrainModelResult dataclass.""" + +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +from model_manager.utils.models.train_model_params import TrainModelParams +from model_manager.utils.models.train_model_result import TrainModelResult + + +@pytest.fixture +def sample_params(): + """Create sample TrainModelParams for testing.""" + return TrainModelParams( + variable_columns=['var1', 'var2'], + lag_train=5, + lag_val=3, + target_variable='target', + rem_static_win=True, + low_lim={'var1': 0.0, 'var2': 0.0}, + upp_lim={'var1': 100.0, 'var2': 100.0}, + window=10, + use_scaler=True, + include_ar=False, + bucket_name='test-bucket', + file_name='test-file.csv', + line_separator='\n', + decimal_separator='.', + train_size=80, + shuffle=True, + experiment_run_id=123, + experiment_name='test-experiment', + experiment_description='Test experiment description', + removed_intervals=[], + ) + + +@pytest.fixture +def sample_dataframes(): + """Create sample DataFrames for testing.""" + X_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]}) + X_test = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10]}) + y_train = pd.DataFrame({'target': [10, 20, 30]}) + y_test = pd.DataFrame({'target': [40, 50]}) + return X_train, X_test, y_train, y_test + + +def test_train_model_result_creation(sample_params, sample_dataframes): + """Test creating TrainModelResult with required fields.""" + X_train, X_test, y_train, y_test = sample_dataframes + process_data = MagicMock() + regr = MagicMock() + scaler_dict = {'var1': {'min': 0, 'max': 100}} + + result = TrainModelResult( + params=sample_params, + process_data=process_data, + X_train=X_train, + X_test=X_test, + y_train=y_train, + y_test=y_test, + regr=regr, + scaler_dict=scaler_dict, + ) + + assert result.params == sample_params + assert result.process_data == process_data + assert result.X_train.equals(X_train) + assert result.X_test.equals(X_test) + assert result.y_train.equals(y_train) + assert result.y_test.equals(y_test) + assert result.regr == regr + assert result.scaler_dict == scaler_dict + + +def test_train_model_result_optional_fields_default_none(sample_params, sample_dataframes): + """Test that optional fields default to None.""" + X_train, X_test, y_train, y_test = sample_dataframes + + result = TrainModelResult( + params=sample_params, + process_data=MagicMock(), + X_train=X_train, + X_test=X_test, + y_train=y_train, + y_test=y_test, + regr=MagicMock(), + scaler_dict={}, + ) + + assert result.y_pred is None + assert result.mse_val is None + assert result.mae_val is None + assert result.r2_val is None + assert result.run_name is None + assert result.report_path is None + assert result.train_data_path is None + assert result.test_data_path is None + assert result.run_dir is None + + +def test_train_model_result_with_metrics(sample_params, sample_dataframes): + """Test TrainModelResult with metrics populated.""" + X_train, X_test, y_train, y_test = sample_dataframes + y_pred = pd.Series([41, 49]) + + result = TrainModelResult( + params=sample_params, + process_data=MagicMock(), + X_train=X_train, + X_test=X_test, + y_train=y_train, + y_test=y_test, + regr=MagicMock(), + scaler_dict={}, + y_pred=y_pred, + mse_val=1.5, + mae_val=1.2, + r2_val=0.95, + ) + + assert result.y_pred.equals(y_pred) + assert result.mse_val == 1.5 + assert result.mae_val == 1.2 + assert result.r2_val == 0.95 + + +def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes): + """Test TrainModelResult with artifact paths populated.""" + X_train, X_test, y_train, y_test = sample_dataframes + + result = TrainModelResult( + params=sample_params, + process_data=MagicMock(), + X_train=X_train, + X_test=X_test, + y_train=y_train, + y_test=y_test, + regr=MagicMock(), + scaler_dict={}, + run_name='test-experiment-1', + report_path='/path/to/report.html', + train_data_path='/path/to/train_data.csv', + test_data_path='/path/to/test_data.csv', + run_dir='/path/to/run_dir', + ) + + assert result.run_name == 'test-experiment-1' + assert result.report_path == '/path/to/report.html' + assert result.train_data_path == '/path/to/train_data.csv' + assert result.test_data_path == '/path/to/test_data.csv' + assert result.run_dir == '/path/to/run_dir' + + +def test_train_model_result_is_dataclass(sample_params, sample_dataframes): + """Test that TrainModelResult is a dataclass.""" + X_train, X_test, y_train, y_test = sample_dataframes + + result = TrainModelResult( + params=sample_params, + process_data=MagicMock(), + X_train=X_train, + X_test=X_test, + y_train=y_train, + y_test=y_test, + regr=MagicMock(), + scaler_dict={}, + ) + + # Dataclasses have __dataclass_fields__ attribute + assert hasattr(result, '__dataclass_fields__') + assert 'params' in result.__dataclass_fields__ + assert 'process_data' in result.__dataclass_fields__ + assert 'X_train' in result.__dataclass_fields__ + + +def test_train_model_result_field_count(): + """Test that TrainModelResult has exactly 17 fields.""" + from dataclasses import fields + + result_fields = fields(TrainModelResult) + assert len(result_fields) == 17 + + field_names = {f.name for f in result_fields} + expected_fields = { + 'params', + 'process_data', + 'X_train', + 'X_test', + 'y_train', + 'y_test', + 'regr', + 'scaler_dict', + 'y_pred', + 'mse_val', + 'mae_val', + 'r2_val', + 'run_name', + 'report_path', + 'train_data_path', + 'test_data_path', + 'run_dir', + } + assert field_names == expected_fields + + +def test_train_model_result_complete_workflow(sample_params, sample_dataframes): + """Test TrainModelResult through a complete workflow simulation.""" + X_train, X_test, y_train, y_test = sample_dataframes + + # Step 1: Create result after training + result = TrainModelResult( + params=sample_params, + process_data=MagicMock(), + X_train=X_train, + X_test=X_test, + y_train=y_train, + y_test=y_test, + regr=MagicMock(), + scaler_dict={'var1': {'min': 0, 'max': 100}}, + ) + + # Step 2: Add predictions and metrics + result.y_pred = pd.Series([41, 49]) + result.mse_val = 1.5 + result.mae_val = 1.2 + result.r2_val = 0.95 + + # Step 3: Add artifact paths + result.run_name = 'test-experiment-1' + result.report_path = '/path/to/report.html' + result.train_data_path = '/path/to/train_data.csv' + result.test_data_path = '/path/to/test_data.csv' + result.run_dir = '/path/to/run_dir' + + # Verify all fields are populated + assert result.y_pred is not None + assert result.mse_val == 1.5 + assert result.mae_val == 1.2 + assert result.r2_val == 0.95 + assert result.run_name == 'test-experiment-1' + assert result.report_path == '/path/to/report.html' + assert result.train_data_path == '/path/to/train_data.csv' + assert result.test_data_path == '/path/to/test_data.csv' + assert result.run_dir == '/path/to/run_dir'