SIENTIAPDE-1241: Added remaining trainig_repository tests and some model_repository tests

This commit is contained in:
Kou-Kinoshita
2025-10-29 11:41:36 -03:00
parent 4fdefd8263
commit 61364671db
2 changed files with 450 additions and 4 deletions

View File

@@ -199,3 +199,291 @@ def test_get_reports_directory(mock_model_serving_class, mock_logger):
assert result.endswith(os.path.join('model_manager', 'reports'))
assert os.path.isabs(result)
@patch('model_manager.utils.repository.model_repository.ModelServing')
def test_init_artifacts_data_success(mock_model_serving_class, mock_logger, mock_train_result):
"""Test _init_artifacts_data successfully prepares data."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
reference_data, current_data = repo._init_artifacts_data(mock_train_result)
# Check reference data
assert 'target' in reference_data.columns
assert 'prediction' in reference_data.columns
assert len(reference_data) == 3
# Check current data
assert 'target' in current_data.columns
assert 'prediction' in current_data.columns
assert len(current_data) == 3
@patch('model_manager.utils.repository.model_repository.ModelServing')
def test_init_artifacts_data_empty_x_train(mock_model_serving_class, mock_logger, mock_train_result):
"""Test _init_artifacts_data raises ValueError when x_train is empty."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
mock_train_result.x_train = pd.DataFrame()
with pytest.raises(ValueError, match='Training features .* are empty'):
repo._init_artifacts_data(mock_train_result)
@patch('model_manager.utils.repository.model_repository.ModelServing')
def test_init_artifacts_data_empty_y_train(mock_model_serving_class, mock_logger, mock_train_result):
"""Test _init_artifacts_data raises ValueError when y_train is empty."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
mock_train_result.y_train = pd.Series(dtype=float)
with pytest.raises(ValueError, match='Training target .* is empty'):
repo._init_artifacts_data(mock_train_result)
@patch('model_manager.utils.repository.model_repository.ModelServing')
def test_init_artifacts_data_none_y_pred(mock_model_serving_class, mock_logger, mock_train_result):
"""Test _init_artifacts_data raises ValueError when y_pred is None."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
mock_train_result.y_pred = None
with pytest.raises(ValueError, match='Test predictions .* are None'):
repo._init_artifacts_data(mock_train_result)
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.datetime')
@patch('model_manager.utils.repository.model_repository.makedirs')
def test_create_run_directory_success(
mock_makedirs, mock_datetime, mock_model_serving_class, mock_logger
):
"""Test _create_run_directory creates directory successfully."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
mock_datetime.now.return_value.strftime.return_value = '20240101_120000_123456'
result = repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
expected_path = os.path.join('/tmp/reports', 'test_run_20240101_120000_123456') # noqa: S108
assert result == expected_path
mock_makedirs.assert_called_once_with(expected_path, exist_ok=True)
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.makedirs')
def test_create_run_directory_permission_error(
mock_makedirs, mock_model_serving_class, mock_logger
):
"""Test _create_run_directory raises PermissionError."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
mock_makedirs.side_effect = PermissionError('Permission denied')
with pytest.raises(PermissionError, match='Permission denied when creating directory'):
repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.shutil.copy')
@patch('builtins.open', create=True)
def test_setup_run_directory_success(mock_open, mock_copy, mock_model_serving_class, mock_logger):
"""Test _setup_run_directory creates files successfully."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
# Check that empty files were created
assert mock_open.call_count == 3
mock_copy.assert_called_once()
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.json.dump')
@patch('model_manager.utils.repository.model_repository.Reports')
@patch('builtins.open', create=True)
def test_generate_report_with_equation(
mock_open, mock_reports_class, mock_json_dump, mock_model_serving_class, mock_logger, mock_train_result
):
"""Test _generate_report creates equation JSON artifact."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
# Add equation to train result
mock_train_result.equation = {
'target_variable': 'target',
'coefficients': {'var1': 1.5, 'var2': -0.75},
'intercept': 10.5,
'equation_string': 'target = 10.5 + 1.5 * var1 + -0.75 * var2',
'latex_equation': 'target = 10.5 + 1.5 \\cdot var1 + -0.75 \\cdot var2',
'model_type': 'Linear Regression',
}
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
# Mock DataFrame.to_csv to avoid file I/O
with patch.object(pd.DataFrame, 'to_csv'):
result = repo._generate_report(reference_data, current_data, mock_train_result)
# Verify equation path was set
assert result.equation_path == os.path.join(mock_train_result.run_dir, 'model_equation.json')
# Verify JSON was written
mock_json_dump.assert_called()
call_args = mock_json_dump.call_args
assert call_args[0][0] == mock_train_result.equation
assert call_args[1]['indent'] == 2
assert call_args[1]['ensure_ascii'] is False
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.Reports')
@patch('builtins.open', create=True)
def test_generate_report_without_equation(
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
):
"""Test _generate_report works without equation."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
# No equation
mock_train_result.equation = None
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
# Mock DataFrame.to_csv to avoid file I/O
with patch.object(pd.DataFrame, 'to_csv'):
result = repo._generate_report(reference_data, current_data, mock_train_result)
# Verify equation path was not set
assert not hasattr(result, 'equation_path') or result.equation_path is None
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.path.exists')
def test_save_run_with_equation(
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
):
"""Test _save_run logs equation artifact."""
from model_manager.utils.repository.model_repository import ModelRepository
mock_model_serving_instance = MagicMock()
mock_model_serving_class.return_value = mock_model_serving_instance
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
# Set equation path
mock_train_result.equation_path = '/tmp/test_run/model_equation.json' # noqa: S108
# Mock all path.exists calls to return True
mock_exists.return_value = True
repo._save_run(mock_train_result)
# Verify equation artifact was logged
logged_artifacts = [
call[0][0] for call in mock_model_serving_instance.log_artifact.call_args_list
]
assert '/tmp/test_run/model_equation.json' in logged_artifacts # noqa: S108
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.path.exists')
def test_save_run_without_equation(
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
):
"""Test _save_run works without equation."""
from model_manager.utils.repository.model_repository import ModelRepository
mock_model_serving_instance = MagicMock()
mock_model_serving_class.return_value = mock_model_serving_instance
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
# No equation
mock_train_result.equation_path = None
# Mock path.exists to return True for required artifacts
mock_exists.return_value = True
repo._save_run(mock_train_result)
# Verify only 3 artifacts were logged (report, train_data, test_data)
assert mock_model_serving_instance.log_artifact.call_count == 3
@patch('model_manager.utils.repository.model_repository.ModelServing')
@patch('model_manager.utils.repository.model_repository.path.exists')
def test_save_run_missing_report(
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
):
"""Test _save_run raises ValueError when report is missing."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
# Mock report doesn't exist
def exists_side_effect(path):
return not path.endswith('report.html')
mock_exists.side_effect = exists_side_effect
with pytest.raises(ValueError, match='Report file does not exist'):
repo._save_run(mock_train_result)
@patch('model_manager.utils.repository.model_repository.ModelServing')
def test_save_run_none_metrics(mock_model_serving_class, mock_logger, mock_train_result):
"""Test _save_run raises ValueError when metrics are None."""
from model_manager.utils.repository.model_repository import ModelRepository
repo = ModelRepository(
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
)
mock_train_result.mse_val = None
with pytest.raises(ValueError, match='One or more metrics .* are None'):
repo._save_run(mock_train_result)

View File

@@ -30,8 +30,6 @@ def sample_params():
return TrainModelParams(
experiment_run_id=1,
experiment_name='test_experiment',
username='test_user',
model_type='Linear Regression',
target_variable='target',
variable_columns=['var1', 'var2', 'var3'],
lag_train=0,
@@ -105,8 +103,6 @@ class TestExtractModelEquation:
params = TrainModelParams(
experiment_run_id=1,
experiment_name='test',
username='test_user',
model_type='Linear Regression',
target_variable='y',
variable_columns=['x'],
lag_train=0,
@@ -224,6 +220,78 @@ class TestInitDataPreprocessor:
assert preprocessor.lag_transform[col] == 3
class TestInitScalerDict:
"""Tests for _init_scaler_dict method."""
def test_init_scaler_dict_without_scaler(self, training_repo, sample_params):
"""Test scaler dict initialization when scaler is not used."""
from model_manager.sientia.models import DataPreprocessor
sample_params.use_scaler = False
process_data = MagicMock(spec=DataPreprocessor)
result = training_repo._init_scaler_dict(process_data, sample_params)
assert result == {}
def test_init_scaler_dict_with_minmax_scaler(self, training_repo, sample_params):
"""Test scaler dict initialization with MinMaxScaler."""
from model_manager.sientia.models import DataPreprocessor
from sientia_do.operations.normalization import MinMaxScaler
sample_params.use_scaler = True
# Create mock MinMaxScaler
mock_scaler = MagicMock(spec=MinMaxScaler)
mock_scaler.x_min = np.array([0.0, 1.0, 2.0])
mock_scaler.x_max = np.array([10.0, 11.0, 12.0])
mock_scaler.y_min = 0.5
mock_scaler.y_max = 100.5
# Create mock preprocessor that returns the scaler
process_data = MagicMock(spec=DataPreprocessor)
process_data.get_scaler.return_value = mock_scaler
result = training_repo._init_scaler_dict(process_data, sample_params)
# Check feature scalers
assert 'var1' in result
assert 'var2' in result
assert 'var3' in result
assert result['var1'] == {'min': 0.0, 'max': 10.0}
assert result['var2'] == {'min': 1.0, 'max': 11.0}
assert result['var3'] == {'min': 2.0, 'max': 12.0}
# Check target scaler
assert 'target' in result
assert result['target'] == {'min': 0.5, 'max': 100.5}
def test_init_scaler_dict_with_z_scaler(self, training_repo, sample_params):
"""Test scaler dict initialization with Z_Scaler."""
from model_manager.sientia.models import DataPreprocessor
from sientia_do.operations.normalization import Z_Scaler
sample_params.use_scaler = True
# Create mock Z_Scaler
mock_scaler = MagicMock(spec=Z_Scaler)
expected_dict = {
'var1': {'mean': 5.0, 'std': 1.5},
'var2': {'mean': 10.0, 'std': 2.0},
'target': {'mean': 50.0, 'std': 10.0},
}
mock_scaler.create_dict.return_value = expected_dict
# Create mock preprocessor
process_data = MagicMock(spec=DataPreprocessor)
process_data.get_scaler.return_value = mock_scaler
result = training_repo._init_scaler_dict(process_data, sample_params)
assert result == expected_dict
mock_scaler.create_dict.assert_called_once()
class TestAfterTrainCalculation:
"""Tests for after_train_calculation method."""
@@ -307,3 +375,93 @@ class TestAfterTrainCalculation:
for call in mock_logger.info.call_args_list
)
def test_after_train_with_custom_scaler_denormalization(
self, training_repo, sample_params, sample_linear_model
):
"""Test denormalization with custom scaler that has denormalize methods."""
sample_params.use_scaler = True
# Create mock data
x_train = pd.DataFrame(
{'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2]
)
x_test = pd.DataFrame({'var1': [10], 'var2': [11], 'var3': [12]}, index=[3])
y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target')
y_test = pd.Series([400], index=[3], name='target')
# Mock predict
sample_linear_model.predict = MagicMock(return_value=np.array([450.0]))
# Create mock scaler with denormalize methods
mock_scaler = MagicMock()
mock_scaler.denormalize_single_input = MagicMock(side_effect=lambda x, col: x * 2)
mock_scaler.denormalize_predictions = MagicMock(return_value=np.array([900.0]))
# Create mock preprocessor
mock_process_data = MagicMock()
mock_process_data.get_scaler.return_value = mock_scaler
train_result = TrainModelResult(
params=sample_params,
process_data=mock_process_data,
x_train=x_train,
x_test=x_test,
y_train=y_train,
y_test=y_test,
regr=sample_linear_model,
scaler_dict={},
)
result = training_repo.after_train_calculation(sample_params, train_result)
# Verify denormalize methods were called
assert mock_scaler.denormalize_single_input.called
assert mock_scaler.denormalize_predictions.called
assert result.y_pred is not None
def test_after_train_with_sklearn_scaler(
self, training_repo, sample_params, sample_linear_model
):
"""Test denormalization with sklearn StandardScaler."""
sample_params.use_scaler = True
# Create mock data
x_train = pd.DataFrame(
{'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2]
)
x_test = pd.DataFrame({'var1': [10], 'var2': [11], 'var3': [12]}, index=[3])
y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target')
y_test = pd.Series([400], index=[3], name='target')
# Mock predict
sample_linear_model.predict = MagicMock(return_value=np.array([450.0]))
# Create mock sklearn scaler (without denormalize methods)
mock_scaler = MagicMock()
# Remove denormalize methods to trigger sklearn path
if hasattr(mock_scaler, 'denormalize_single_input'):
delattr(mock_scaler, 'denormalize_single_input')
mock_scaler.inverse_transform = MagicMock(side_effect=lambda x: x * 2)
# Create mock preprocessor with feature_names_order
mock_process_data = MagicMock()
mock_process_data.get_scaler.return_value = mock_scaler
mock_process_data.feature_names_order = ['var1', 'var2', 'var3']
train_result = TrainModelResult(
params=sample_params,
process_data=mock_process_data,
x_train=x_train,
x_test=x_test,
y_train=y_train,
y_test=y_test,
regr=sample_linear_model,
scaler_dict={},
)
result = training_repo.after_train_calculation(sample_params, train_result)
# Verify inverse_transform was called
assert mock_scaler.inverse_transform.called
assert result.y_pred is not None