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
0
tests/laborious/utils/models/__init__.py
Normal file
0
tests/laborious/utils/models/__init__.py
Normal file
79
tests/laborious/utils/models/test_experiment_status.py
Normal file
79
tests/laborious/utils/models/test_experiment_status.py
Normal 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
|
||||
51
tests/laborious/utils/models/test_init.py
Normal file
51
tests/laborious/utils/models/test_init.py
Normal 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))
|
||||
295
tests/laborious/utils/models/test_train_model_params.py
Normal file
295
tests/laborious/utils/models/test_train_model_params.py
Normal 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)
|
||||
246
tests/laborious/utils/models/test_train_model_result.py
Normal file
246
tests/laborious/utils/models/test_train_model_result.py
Normal 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'
|
||||
Reference in New Issue
Block a user