gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1249: Refactor TrainModelParams to use dataclass and add from_dict method for validation, remove no-cache-dir from pip install in quality gate workflow, and rename X_train/X_test to x_train/x_test in TrainModelResult.

Browse Source
This commit is contained in:
Bruno Domingues
2025-10-06 19:11:02 -03:00
parent d8583cdae7
commit 04c94efd98
5 changed files with 153 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.github/workflows/quality-gate.yml vendored
View File

@@ -208,8 +208,8 @@ jobs:
- name: 📦 Install Dependencies - name: 📦 Install Dependencies
run: | run: |
python -m pip install --upgrade pip python -m pip install --upgrade pip
pip install --no-cache-dir -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }} pip install -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }}
pip install --no-cache-dir -r requirements-dev.txt pip install -r requirements-dev.txt
- name: 📝 Code Formatting Check (Ruff) - name: 📝 Code Formatting Check (Ruff)
run: | run: |

170
model_manager/utils/models/train_model_params.py
View File

@@ -1,6 +1,8 @@
from dataclasses import dataclass
from typing import Any from typing import Any
@dataclass
class TrainModelParams: class TrainModelParams:
""" """
Parameters for machine learning model training. Parameters for machine learning model training.
@@ -10,14 +12,17 @@ class TrainModelParams:
tracking information. All parameters are validated upon initialization to ensure tracking information. All parameters are validated upon initialization to ensure
data integrity and prevent runtime errors. data integrity and prevent runtime errors.
Use the `from_dict()` class method to create instances from dictionaries with
automatic validation of all fields.
Attributes: Attributes:
variable_columns (List[str]): List of variable column names to use as features. 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_train (int): Number of lags to apply during training phase.
lag_val (int): Number of lags to apply during validation phase. lag_val (int): Number of lags to apply during validation phase.
target_variable (str): Name of the target variable to predict. target_variable (str): Name of the target variable to predict.
rem_static_win (bool): Whether to remove static windows from data. rem_static_win (bool): Whether to remove static windows from data.
low_lim (Dict[str, float]): Dictionary of lower limits for each variable. low_lim (dict[str, float]): Dictionary of lower limits for each variable.
upp_lim (Dict[str, float]): Dictionary of upper limits for each variable. upp_lim (dict[str, float]): Dictionary of upper limits for each variable.
window (int): Window size for rolling operations. window (int): Window size for rolling operations.
use_scaler (bool): Whether to use a scaler for data normalization. use_scaler (bool): Whether to use a scaler for data normalization.
include_ar (bool): Whether to include autoregressive variables. include_ar (bool): Whether to include autoregressive variables.
@@ -33,90 +38,92 @@ class TrainModelParams:
removed_intervals (list): List of time intervals to remove from the data. removed_intervals (list): List of time intervals to remove from the data.
""" """
def __init__( variable_columns: list[str]
self, lag_train: int
variable_columns: Any | None, lag_val: int
lag_train: Any | None, target_variable: str
lag_val: Any | None, rem_static_win: bool
target_variable: Any | None, low_lim: dict[str, float]
rem_static_win: Any | None, upp_lim: dict[str, float]
low_lim: Any | None, window: int
upp_lim: Any | None, use_scaler: bool
window: Any | None, include_ar: bool
use_scaler: Any | None, bucket_name: str
include_ar: Any | None, file_name: str
bucket_name: Any | None, line_separator: str
file_name: Any | None, decimal_separator: str
line_separator: Any | None, train_size: int
decimal_separator: Any | None, shuffle: bool
train_size: Any | None, experiment_run_id: int
shuffle: Any | None, experiment_name: str
experiment_run_id: Any | None, experiment_description: str
experiment_name: Any | None, removed_intervals: list
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 @classmethod
have the correct types. It raises descriptive errors if validation fails, def from_dict(cls, data: dict[str, Any]) -> 'TrainModelParams':
helping to catch configuration issues early in the pipeline. """
Create TrainModelParams from dictionary with validation.
This factory method creates a TrainModelParams instance from a dictionary,
applying validation to ensure all required fields are present and have
the correct types. This is the recommended way to create instances from
workflow input data.
Args: Args:
variable_columns (list | None): List of variable column names. data: Dictionary containing training parameters with keys matching
lag_train (int | None): Number of lags to apply during training. the attribute names (variable_columns, lag_train, etc.)
lag_val (int | None): Number of lags to apply during validation.
target_variable (str | None): Name of the target variable. Returns:
rem_static_win (bool | None): Whether to remove static windows. TrainModelParams: Validated instance with all fields populated
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: Raises:
ValueError: If any required parameter is None. ValueError: If any required field is missing or None
TypeError: If any parameter is not of the expected type. TypeError: If any field has an incorrect type
KeyError: If any required key is missing from the dictionary
Example:
>>> input_data = {
... 'variable_columns': ['var1', 'var2'],
... 'lag_train': 5,
... # ... other fields
... }
>>> params = TrainModelParams.from_dict(input_data)
""" """
self.variable_columns: list[str] = self._check_none( return cls(
variable_columns, list, 'variable_columns' variable_columns=cls._check_none(
) data.get('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') lag_train=cls._check_none(data.get('lag_train'), int, 'lag_train'),
self.target_variable: str = self._check_none(target_variable, str, 'target_variable') lag_val=cls._check_none(data.get('lag_val'), int, 'lag_val'),
self.rem_static_win: bool = self._check_none(rem_static_win, bool, 'rem_static_win') target_variable=cls._check_none(data.get('target_variable'), str, 'target_variable'),
self.low_lim: dict[str, float] = self._check_none(low_lim, dict, 'low_lim') rem_static_win=cls._check_none(data.get('rem_static_win'), bool, 'rem_static_win'),
self.upp_lim: dict[str, float] = self._check_none(upp_lim, dict, 'upp_lim') low_lim=cls._check_none(data.get('low_lim'), dict, 'low_lim'),
self.window: int = self._check_none(window, int, 'window') upp_lim=cls._check_none(data.get('upp_lim'), dict, 'upp_lim'),
self.use_scaler: bool = self._check_none(use_scaler, bool, 'use_scaler') window=cls._check_none(data.get('window'), int, 'window'),
self.include_ar: bool = self._check_none(include_ar, bool, 'include_ar') use_scaler=cls._check_none(data.get('use_scaler'), bool, 'use_scaler'),
self.bucket_name: str = self._check_none(bucket_name, str, 'bucket_name') include_ar=cls._check_none(data.get('include_ar'), bool, 'include_ar'),
self.file_name: str = self._check_none(file_name, str, 'file_name') bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'),
self.line_separator: str = self._check_none(line_separator, str, 'line_separator') file_name=cls._check_none(data.get('file_name'), str, 'file_name'),
self.decimal_separator: str = self._check_none(decimal_separator, str, 'decimal_separator') line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'),
self.train_size: int = self._check_none(train_size, int, 'train_size') decimal_separator=cls._check_none(
self.shuffle: bool = self._check_none(shuffle, bool, 'shuffle') data.get('decimal_separator'), str, 'decimal_separator'
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') train_size=cls._check_none(data.get('train_size'), int, 'train_size'),
self.experiment_description: str = self._check_none( shuffle=cls._check_none(data.get('shuffle'), bool, 'shuffle'),
experiment_description, str, 'experiment_description' experiment_run_id=cls._check_none(
) data.get('experiment_run_id'), int, 'experiment_run_id'
self.removed_intervals: list = self._check_type( ),
removed_intervals, list, 'removed_intervals' experiment_name=cls._check_none(data.get('experiment_name'), str, 'experiment_name'),
experiment_description=cls._check_none(
data.get('experiment_description'), str, 'experiment_description'
),
removed_intervals=cls._check_type(
data.get('removed_intervals'), list, 'removed_intervals'
),
) )
def _check_none(self, value: Any | None, expected_type: type, field_name: str): @staticmethod
def _check_none(value: Any | None, expected_type: type, field_name: str) -> Any:
""" """
Validate that a value is not None and check its type. Validate that a value is not None and check its type.
@@ -139,9 +146,10 @@ class TrainModelParams:
error = f'{field_name} is required and cannot be None.' error = f'{field_name} is required and cannot be None.'
raise ValueError(error) raise ValueError(error)
return self._check_type(value, expected_type, field_name) return TrainModelParams._check_type(value, expected_type, field_name)
def _check_type(self, value: Any | None, expected_type: type, field_name: str): @staticmethod
def _check_type(value: Any | None, expected_type: type, field_name: str) -> Any:
""" """
Validate that a value matches the expected type. Validate that a value matches the expected type.

26
model_manager/utils/models/train_model_result.py
View File

@@ -19,27 +19,27 @@ class TrainModelResult:
Attributes: Attributes:
params (TrainModelParams): The parameters used to train the model. params (TrainModelParams): The parameters used to train the model.
process_data (DataPreprocessor): The data preprocessor object used to process the input data. process_data (DataPreprocessor): The data preprocessor object used to process the input data.
X_train (pd.DataFrame): The training dataset features. x_train (pd.DataFrame): The training dataset features.
X_test (pd.DataFrame): The testing dataset features. x_test (pd.DataFrame): The testing dataset features.
y_train (pd.DataFrame): The training dataset target values. y_train (pd.DataFrame): The training dataset target values.
y_test (pd.DataFrame): The testing dataset target values. y_test (pd.DataFrame): The testing dataset target values.
regr (LinearRegressionModel): The trained linear regression model. regr (LinearRegressionModel): The trained linear regression model.
scaler_dict (dict): A dictionary containing the scalers used to scale the features and target values. 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. y_pred (pd.Series | None): 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. mse_val (float | None): 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. mae_val (float | None): 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. r2_val (float | None): The R-squared (R²) value of the predictions. Default is None.
run_name (Optional[str]): The name of the MLFlow run. Default is None. run_name (str | None): The name of the MLFlow run. Default is None.
report_path (Optional[str]): The path to the generated HTML report file. Default is None. report_path (str | None): 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. train_data_path (str | None): 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. test_data_path (str | None): 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. run_dir (str | None): The path to the run directory containing all artifacts. Default is None.
""" """
params: TrainModelParams params: TrainModelParams
process_data: DataPreprocessor process_data: DataPreprocessor
X_train: pd.DataFrame x_train: pd.DataFrame
X_test: pd.DataFrame x_test: pd.DataFrame
y_train: pd.DataFrame y_train: pd.DataFrame
y_test: pd.DataFrame y_test: pd.DataFrame
regr: LinearRegressionModel regr: LinearRegressionModel

43
tests/laborious/utils/models/test_train_model_params.py
View File

@@ -58,12 +58,21 @@ def test_train_model_params_creation_with_valid_params(valid_params_dict):
assert params.removed_intervals == [] assert params.removed_intervals == []
def test_train_model_params_from_dict_creation(valid_params_dict):
"""Test creating TrainModelParams using from_dict method."""
params = TrainModelParams.from_dict(valid_params_dict)
assert params.variable_columns == ['var1', 'var2', 'var3']
assert params.lag_train == 5
assert params.experiment_run_id == 123
def test_train_model_params_variable_columns_none_raises_error(valid_params_dict): def test_train_model_params_variable_columns_none_raises_error(valid_params_dict):
"""Test that None variable_columns raises ValueError.""" """Test that None variable_columns raises ValueError."""
valid_params_dict['variable_columns'] = None valid_params_dict['variable_columns'] = None
with pytest.raises(ValueError, match='variable_columns is required and cannot be None'): with pytest.raises(ValueError, match='variable_columns is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_variable_columns_wrong_type_raises_error(valid_params_dict): def test_train_model_params_variable_columns_wrong_type_raises_error(valid_params_dict):
@@ -71,7 +80,7 @@ def test_train_model_params_variable_columns_wrong_type_raises_error(valid_param
valid_params_dict['variable_columns'] = 'not a list' valid_params_dict['variable_columns'] = 'not a list'
with pytest.raises(TypeError, match='variable_columns must be of type list'): with pytest.raises(TypeError, match='variable_columns must be of type list'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_lag_train_none_raises_error(valid_params_dict): def test_train_model_params_lag_train_none_raises_error(valid_params_dict):
@@ -79,7 +88,7 @@ def test_train_model_params_lag_train_none_raises_error(valid_params_dict):
valid_params_dict['lag_train'] = None valid_params_dict['lag_train'] = None
with pytest.raises(ValueError, match='lag_train is required and cannot be None'): with pytest.raises(ValueError, match='lag_train is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_lag_train_wrong_type_raises_error(valid_params_dict): def test_train_model_params_lag_train_wrong_type_raises_error(valid_params_dict):
@@ -87,7 +96,7 @@ def test_train_model_params_lag_train_wrong_type_raises_error(valid_params_dict)
valid_params_dict['lag_train'] = '5' valid_params_dict['lag_train'] = '5'
with pytest.raises(TypeError, match='lag_train must be of type int'): with pytest.raises(TypeError, match='lag_train must be of type int'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_target_variable_none_raises_error(valid_params_dict): def test_train_model_params_target_variable_none_raises_error(valid_params_dict):
@@ -95,7 +104,7 @@ def test_train_model_params_target_variable_none_raises_error(valid_params_dict)
valid_params_dict['target_variable'] = None valid_params_dict['target_variable'] = None
with pytest.raises(ValueError, match='target_variable is required and cannot be None'): with pytest.raises(ValueError, match='target_variable is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_target_variable_wrong_type_raises_error(valid_params_dict): def test_train_model_params_target_variable_wrong_type_raises_error(valid_params_dict):
@@ -103,7 +112,7 @@ def test_train_model_params_target_variable_wrong_type_raises_error(valid_params
valid_params_dict['target_variable'] = 123 valid_params_dict['target_variable'] = 123
with pytest.raises(TypeError, match='target_variable must be of type str'): with pytest.raises(TypeError, match='target_variable must be of type str'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_boolean_fields(valid_params_dict): def test_train_model_params_boolean_fields(valid_params_dict):
@@ -111,11 +120,11 @@ def test_train_model_params_boolean_fields(valid_params_dict):
# Test rem_static_win # Test rem_static_win
valid_params_dict['rem_static_win'] = None valid_params_dict['rem_static_win'] = None
with pytest.raises(ValueError, match='rem_static_win is required and cannot be None'): with pytest.raises(ValueError, match='rem_static_win is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
valid_params_dict['rem_static_win'] = 'true' valid_params_dict['rem_static_win'] = 'true'
with pytest.raises(TypeError, match='rem_static_win must be of type bool'): with pytest.raises(TypeError, match='rem_static_win must be of type bool'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_dict_fields(valid_params_dict): def test_train_model_params_dict_fields(valid_params_dict):
@@ -123,12 +132,12 @@ def test_train_model_params_dict_fields(valid_params_dict):
# Test low_lim # Test low_lim
valid_params_dict['low_lim'] = None valid_params_dict['low_lim'] = None
with pytest.raises(ValueError, match='low_lim is required and cannot be None'): with pytest.raises(ValueError, match='low_lim is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
valid_params_dict['low_lim'] = {'var1': 0.0} valid_params_dict['low_lim'] = {'var1': 0.0}
valid_params_dict['upp_lim'] = 'not a dict' valid_params_dict['upp_lim'] = 'not a dict'
with pytest.raises(TypeError, match='upp_lim must be of type dict'): with pytest.raises(TypeError, match='upp_lim must be of type dict'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_bucket_name_none_raises_error(valid_params_dict): def test_train_model_params_bucket_name_none_raises_error(valid_params_dict):
@@ -136,7 +145,7 @@ def test_train_model_params_bucket_name_none_raises_error(valid_params_dict):
valid_params_dict['bucket_name'] = None valid_params_dict['bucket_name'] = None
with pytest.raises(ValueError, match='bucket_name is required and cannot be None'): with pytest.raises(ValueError, match='bucket_name is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_file_name_none_raises_error(valid_params_dict): def test_train_model_params_file_name_none_raises_error(valid_params_dict):
@@ -144,7 +153,7 @@ def test_train_model_params_file_name_none_raises_error(valid_params_dict):
valid_params_dict['file_name'] = None valid_params_dict['file_name'] = None
with pytest.raises(ValueError, match='file_name is required and cannot be None'): with pytest.raises(ValueError, match='file_name is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_experiment_run_id_none_raises_error(valid_params_dict): def test_train_model_params_experiment_run_id_none_raises_error(valid_params_dict):
@@ -152,7 +161,7 @@ def test_train_model_params_experiment_run_id_none_raises_error(valid_params_dic
valid_params_dict['experiment_run_id'] = None valid_params_dict['experiment_run_id'] = None
with pytest.raises(ValueError, match='experiment_run_id is required and cannot be None'): with pytest.raises(ValueError, match='experiment_run_id is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_experiment_name_none_raises_error(valid_params_dict): def test_train_model_params_experiment_name_none_raises_error(valid_params_dict):
@@ -160,14 +169,14 @@ def test_train_model_params_experiment_name_none_raises_error(valid_params_dict)
valid_params_dict['experiment_name'] = None valid_params_dict['experiment_name'] = None
with pytest.raises(ValueError, match='experiment_name is required and cannot be None'): with pytest.raises(ValueError, match='experiment_name is required and cannot be None'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_removed_intervals_can_be_none(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).""" """Test that removed_intervals can be None (uses _check_type not _check_none)."""
valid_params_dict['removed_intervals'] = None valid_params_dict['removed_intervals'] = None
params = TrainModelParams(**valid_params_dict) params = TrainModelParams.from_dict(valid_params_dict)
assert params.removed_intervals is None assert params.removed_intervals is None
@@ -176,7 +185,7 @@ def test_train_model_params_removed_intervals_wrong_type_raises_error(valid_para
valid_params_dict['removed_intervals'] = 'not a list' valid_params_dict['removed_intervals'] = 'not a list'
with pytest.raises(TypeError, match='removed_intervals must be of type list'): with pytest.raises(TypeError, match='removed_intervals must be of type list'):
TrainModelParams(**valid_params_dict) TrainModelParams.from_dict(valid_params_dict)
def test_train_model_params_removed_intervals_with_values(valid_params_dict): def test_train_model_params_removed_intervals_with_values(valid_params_dict):
@@ -186,7 +195,7 @@ def test_train_model_params_removed_intervals_with_values(valid_params_dict):
('2023-02-01', '2023-02-05'), ('2023-02-01', '2023-02-05'),
] ]
params = TrainModelParams(**valid_params_dict) params = TrainModelParams.from_dict(valid_params_dict)
assert len(params.removed_intervals) == 2 assert len(params.removed_intervals) == 2
assert params.removed_intervals[0] == ('2023-01-01', '2023-01-10') assert params.removed_intervals[0] == ('2023-01-01', '2023-01-10')

46
tests/laborious/utils/models/test_train_model_result.py
View File

@@ -48,7 +48,7 @@ def sample_dataframes():
def test_train_model_result_creation(sample_params, sample_dataframes): def test_train_model_result_creation(sample_params, sample_dataframes):
"""Test creating TrainModelResult with required fields.""" """Test creating TrainModelResult with required fields."""
X_train, X_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
process_data = MagicMock() process_data = MagicMock()
regr = MagicMock() regr = MagicMock()
scaler_dict = {'var1': {'min': 0, 'max': 100}} scaler_dict = {'var1': {'min': 0, 'max': 100}}
@@ -56,8 +56,8 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
process_data=process_data, process_data=process_data,
X_train=X_train, x_train=x_train,
X_test=X_test, x_test=x_test,
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=regr, regr=regr,
@@ -66,8 +66,8 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
assert result.params == sample_params assert result.params == sample_params
assert result.process_data == process_data assert result.process_data == process_data
assert result.X_train.equals(X_train) assert result.x_train.equals(x_train)
assert result.X_test.equals(X_test) assert result.x_test.equals(x_test)
assert result.y_train.equals(y_train) assert result.y_train.equals(y_train)
assert result.y_test.equals(y_test) assert result.y_test.equals(y_test)
assert result.regr == regr assert result.regr == regr
@@ -76,13 +76,13 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
def test_train_model_result_optional_fields_default_none(sample_params, sample_dataframes): def test_train_model_result_optional_fields_default_none(sample_params, sample_dataframes):
"""Test that optional fields default to None.""" """Test that optional fields default to None."""
X_train, X_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
process_data=MagicMock(), process_data=MagicMock(),
X_train=X_train, x_train=x_train,
X_test=X_test, x_test=x_test,
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
@@ -102,14 +102,14 @@ def test_train_model_result_optional_fields_default_none(sample_params, sample_d
def test_train_model_result_with_metrics(sample_params, sample_dataframes): def test_train_model_result_with_metrics(sample_params, sample_dataframes):
"""Test TrainModelResult with metrics populated.""" """Test TrainModelResult with metrics populated."""
X_train, X_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
y_pred = pd.Series([41, 49]) y_pred = pd.Series([41, 49])
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
process_data=MagicMock(), process_data=MagicMock(),
X_train=X_train, x_train=x_train,
X_test=X_test, x_test=x_test,
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
@@ -128,13 +128,13 @@ def test_train_model_result_with_metrics(sample_params, sample_dataframes):
def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes): def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes):
"""Test TrainModelResult with artifact paths populated.""" """Test TrainModelResult with artifact paths populated."""
X_train, X_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
process_data=MagicMock(), process_data=MagicMock(),
X_train=X_train, x_train=x_train,
X_test=X_test, x_test=x_test,
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
@@ -155,13 +155,13 @@ def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes
def test_train_model_result_is_dataclass(sample_params, sample_dataframes): def test_train_model_result_is_dataclass(sample_params, sample_dataframes):
"""Test that TrainModelResult is a dataclass.""" """Test that TrainModelResult is a dataclass."""
X_train, X_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
process_data=MagicMock(), process_data=MagicMock(),
X_train=X_train, x_train=x_train,
X_test=X_test, x_test=x_test,
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
@@ -172,7 +172,7 @@ def test_train_model_result_is_dataclass(sample_params, sample_dataframes):
assert hasattr(result, '__dataclass_fields__') assert hasattr(result, '__dataclass_fields__')
assert 'params' in result.__dataclass_fields__ assert 'params' in result.__dataclass_fields__
assert 'process_data' in result.__dataclass_fields__ assert 'process_data' in result.__dataclass_fields__
assert 'X_train' in result.__dataclass_fields__ assert 'x_train' in result.__dataclass_fields__
def test_train_model_result_field_count(): def test_train_model_result_field_count():
@@ -186,8 +186,8 @@ def test_train_model_result_field_count():
expected_fields = { expected_fields = {
'params', 'params',
'process_data', 'process_data',
'X_train', 'x_train',
'X_test', 'x_test',
'y_train', 'y_train',
'y_test', 'y_test',
'regr', 'regr',
@@ -207,14 +207,14 @@ def test_train_model_result_field_count():
def test_train_model_result_complete_workflow(sample_params, sample_dataframes): def test_train_model_result_complete_workflow(sample_params, sample_dataframes):
"""Test TrainModelResult through a complete workflow simulation.""" """Test TrainModelResult through a complete workflow simulation."""
X_train, X_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
# Step 1: Create result after training # Step 1: Create result after training
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
process_data=MagicMock(), process_data=MagicMock(),
X_train=X_train, x_train=x_train,
X_test=X_test, x_test=x_test,
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),