SIENTIAPDE-1241: Added remaining training repository tests
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
"""Unit tests for TrainingRepository."""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
@@ -466,3 +467,263 @@ class TestAfterTrainCalculation:
|
||||
# Verify inverse_transform was called
|
||||
assert mock_scaler.inverse_transform.called
|
||||
assert result.y_pred is not None
|
||||
|
||||
|
||||
class TestTrain:
|
||||
"""Tests for train method."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_csv_data(self):
|
||||
"""Create sample CSV data in BytesIO."""
|
||||
csv_content = """var1,var2,var3,target
|
||||
1.0,2.0,3.0,10.0
|
||||
2.0,3.0,4.0,15.0
|
||||
3.0,4.0,5.0,20.0
|
||||
4.0,5.0,6.0,25.0
|
||||
5.0,6.0,7.0,30.0
|
||||
6.0,7.0,8.0,35.0
|
||||
7.0,8.0,9.0,40.0
|
||||
8.0,9.0,10.0,45.0
|
||||
9.0,10.0,11.0,50.0
|
||||
10.0,11.0,12.0,55.0
|
||||
"""
|
||||
return BytesIO(csv_content.encode('utf-8'))
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_basic_workflow(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test basic training workflow."""
|
||||
# Mock load_data to return a DataFrame
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
# Mock split_train_test to return train/test splits
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
result = training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
# Verify result structure
|
||||
assert result is not None
|
||||
assert isinstance(result, TrainModelResult)
|
||||
assert result.regr is not None
|
||||
assert result.x_train is not None
|
||||
assert result.x_test is not None
|
||||
assert result.y_train is not None
|
||||
assert result.y_test is not None
|
||||
assert result.process_data is not None
|
||||
assert result.scaler_dict is not None
|
||||
|
||||
# Verify load_data was called correctly
|
||||
mock_load_data.assert_called_once_with(
|
||||
sample_csv_data, sample_params.line_separator, sample_params.decimal_separator
|
||||
)
|
||||
|
||||
# Verify split was called
|
||||
assert mock_split_train_test.called
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_with_scaler(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test training with scaler enabled."""
|
||||
sample_params.use_scaler = True
|
||||
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
result = training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
assert result is not None
|
||||
assert result.scaler_dict is not None
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_with_shuffle_enabled(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test training with shuffle enabled."""
|
||||
sample_params.shuffle = True
|
||||
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
result = training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
# Verify split was called with shuffle=True
|
||||
call_kwargs = mock_split_train_test.call_args[1]
|
||||
assert call_kwargs['shuffle'] is True
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_with_different_train_size(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test training with different train size."""
|
||||
sample_params.train_size = 70
|
||||
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
result = training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
# Verify split was called with train_size=0.7
|
||||
call_kwargs = mock_split_train_test.call_args[1]
|
||||
assert call_kwargs['train_size'] == 0.7
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_raises_on_empty_data_after_transform(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test that ValueError is raised when transformed data is empty."""
|
||||
# Mock load_data to return empty DataFrame
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [],
|
||||
'var2': [],
|
||||
'var3': [],
|
||||
'target': [],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
with pytest.raises(ValueError, match='Data view is empty after transformation'):
|
||||
training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_logs_success(
|
||||
self,
|
||||
mock_load_data,
|
||||
mock_split_train_test,
|
||||
training_repo,
|
||||
mock_logger,
|
||||
sample_params,
|
||||
sample_csv_data,
|
||||
):
|
||||
"""Test that training success is logged."""
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
# Verify success was logged
|
||||
mock_logger.info.assert_called()
|
||||
assert any(
|
||||
'Model trained successfully' in str(call) for call in mock_logger.info.call_args_list
|
||||
)
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_with_custom_separators(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test training with custom line and decimal separators."""
|
||||
sample_params.line_separator = ';'
|
||||
sample_params.decimal_separator = ','
|
||||
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
# Verify load_data was called with custom separators
|
||||
mock_load_data.assert_called_once_with(sample_csv_data, ';', ',')
|
||||
|
||||
@patch('model_manager.utils.repository.training_repository.split_train_test')
|
||||
@patch('model_manager.utils.repository.training_repository.load_data')
|
||||
def test_train_result_contains_all_fields(
|
||||
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
|
||||
):
|
||||
"""Test that TrainModelResult contains all expected fields."""
|
||||
mock_df = pd.DataFrame({
|
||||
'var1': [1, 2, 3, 4, 5],
|
||||
'var2': [2, 3, 4, 5, 6],
|
||||
'var3': [3, 4, 5, 6, 7],
|
||||
'target': [10, 15, 20, 25, 30],
|
||||
})
|
||||
mock_load_data.return_value = mock_df
|
||||
|
||||
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
|
||||
x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
|
||||
y_train = pd.Series([10, 15, 20], name='target')
|
||||
y_test = pd.Series([25, 30], name='target')
|
||||
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
|
||||
|
||||
result = training_repo.train(sample_csv_data, sample_params)
|
||||
|
||||
# Verify all expected fields are present
|
||||
assert hasattr(result, 'params')
|
||||
assert hasattr(result, 'process_data')
|
||||
assert hasattr(result, 'x_train')
|
||||
assert hasattr(result, 'x_test')
|
||||
assert hasattr(result, 'y_train')
|
||||
assert hasattr(result, 'y_test')
|
||||
assert hasattr(result, 'regr')
|
||||
assert hasattr(result, 'scaler_dict')
|
||||
assert result.params == sample_params
|
||||
|
||||
Reference in New Issue
Block a user