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.

This commit is contained in:
Bruno Domingues
2025-10-06 16:49:18 -03:00
parent 18ccfed78a
commit d8583cdae7
9 changed files with 940 additions and 0 deletions

View File

@@ -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',
]

View File

@@ -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'

View File

@@ -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

View File

@@ -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

View File

View File

@@ -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

View File

@@ -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))

View File

@@ -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)

View File

@@ -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'