From 17a551d6f909b0ac21002eab069916dd1ab6e030 Mon Sep 17 00:00:00 2001 From: Kou-Kinoshita Date: Thu, 30 Oct 2025 10:29:51 -0300 Subject: [PATCH] SIENTIAPDE-1241: Added remaining training repository tests --- .../repository/test_training_repository.py | 263 +++++++++++++++++- 1 file changed, 262 insertions(+), 1 deletion(-) diff --git a/tests/utils/repository/test_training_repository.py b/tests/utils/repository/test_training_repository.py index 7707bda..22d739e 100644 --- a/tests/utils/repository/test_training_repository.py +++ b/tests/utils/repository/test_training_repository.py @@ -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