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

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'