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'