SIENTIAPDE-1241: Added training repository tests

2025-10-29 11:32:24 -03:00
parent 8103f530e2
commit 4fdefd8263
1 changed files with 309 additions and 0 deletions
--- a/tests/utils/repository/test_training_repository.py
+++ b/tests/utils/repository/test_training_repository.py
@@ -0,0 +1,309 @@
+"""Unit tests for TrainingRepository."""
+
+from unittest.mock import MagicMock, Mock, patch
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from model_manager.sientia.models import LinearRegressionModel
+from model_manager.utils.models.train_model_params import TrainModelParams
+from model_manager.utils.models.train_model_result import TrainModelResult
+from model_manager.utils.repository.training_repository import TrainingRepository
+
+
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger for testing."""
+    return MagicMock()
+
+
+@pytest.fixture
+def training_repo(mock_logger):
+    """Create TrainingRepository instance with mock logger."""
+    return TrainingRepository(logger=mock_logger)
+
+
+@pytest.fixture
+def sample_params():
+    """Create sample TrainModelParams for testing."""
+    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,
+        lag_val=0,
+        rem_static_win=False,
+        low_lim={},
+        upp_lim={},
+        window=0,
+        use_scaler=False,
+        include_ar=False,
+        train_size=80,
+        shuffle=True,
+        bucket_name='test-bucket',
+        file_name='test.csv',
+        line_separator=',',
+        decimal_separator='.',
+        removed_intervals=[],
+    )
+
+
+@pytest.fixture
+def sample_linear_model():
+    """Create a mock LinearRegressionModel with known coefficients."""
+    mock_model = MagicMock(spec=LinearRegressionModel)
+    mock_sklearn_model = MagicMock()
+    mock_sklearn_model.coef_ = np.array([1.5, -0.75, 2.25])
+    mock_sklearn_model.intercept_ = 10.5
+    mock_model.regr = mock_sklearn_model
+    return mock_model
+
+
+class TestExtractModelEquation:
+    """Tests for _extract_model_equation method."""
+
+    def test_extract_equation_basic(self, training_repo, sample_params, sample_linear_model):
+        """Test basic equation extraction with simple coefficients."""
+        result = training_repo._extract_model_equation(sample_linear_model, sample_params)
+
+        assert result['target_variable'] == 'target'
+        assert result['model_type'] == 'Linear Regression'
+        assert result['intercept'] == 10.5
+
+        assert 'var1' in result['coefficients']
+        assert 'var2' in result['coefficients']
+        assert 'var3' in result['coefficients']
+
+        assert result['coefficients']['var1'] == 1.5
+        assert result['coefficients']['var2'] == -0.75
+        assert result['coefficients']['var3'] == 2.25
+
+    def test_extract_equation_string_format(self, training_repo, sample_params, sample_linear_model):
+        """Test equation string is formatted correctly."""
+        result = training_repo._extract_model_equation(sample_linear_model, sample_params)
+
+        expected_string = (
+            'target = 10.500000 + 1.500000 * var1 + -0.750000 * var2 + 2.250000 * var3'
+        )
+        assert result['equation_string'] == expected_string
+
+    def test_extract_equation_latex_format(self, training_repo, sample_params, sample_linear_model):
+        """Test LaTeX equation is formatted correctly."""
+        result = training_repo._extract_model_equation(sample_linear_model, sample_params)
+
+        expected_latex = (
+            'target = 10.500000 + 1.500000 \\cdot var1 + -0.750000 \\cdot var2 + 2.250000 \\cdot var3'
+        )
+        assert result['latex_equation'] == expected_latex
+
+    def test_extract_equation_single_variable(self, training_repo, sample_linear_model):
+        """Test equation extraction with single variable."""
+        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,
+            lag_val=0,
+            rem_static_win=False,
+            low_lim={},
+            upp_lim={},
+            window=0,
+            use_scaler=False,
+            include_ar=False,
+            train_size=80,
+            shuffle=True,
+            bucket_name='bucket',
+            file_name='file.csv',
+            line_separator=',',
+            decimal_separator='.',
+            removed_intervals=[],
+        )
+
+        # Mock model with single coefficient
+        mock_model = MagicMock(spec=LinearRegressionModel)
+        mock_sklearn_model = MagicMock()
+        mock_sklearn_model.coef_ = np.array([3.14])
+        mock_sklearn_model.intercept_ = 2.71
+        mock_model.regr = mock_sklearn_model
+
+        result = training_repo._extract_model_equation(mock_model, params)
+
+        assert len(result['coefficients']) == 1
+        assert result['coefficients']['x'] == 3.14
+        assert result['intercept'] == 2.71
+        assert 'y = 2.710000 + 3.140000 * x' == result['equation_string']
+
+    def test_extract_equation_zero_coefficients(self, training_repo, sample_params):
+        """Test equation extraction when coefficients are zero."""
+        mock_model = MagicMock(spec=LinearRegressionModel)
+        mock_sklearn_model = MagicMock()
+        mock_sklearn_model.coef_ = np.array([0.0, 0.0, 0.0])
+        mock_sklearn_model.intercept_ = 5.0
+        mock_model.regr = mock_sklearn_model
+
+        result = training_repo._extract_model_equation(mock_model, sample_params)
+
+        assert all(v == 0.0 for v in result['coefficients'].values())
+        assert result['intercept'] == 5.0
+
+    def test_extract_equation_negative_intercept(self, training_repo, sample_params):
+        """Test equation extraction with negative intercept."""
+        mock_model = MagicMock(spec=LinearRegressionModel)
+        mock_sklearn_model = MagicMock()
+        mock_sklearn_model.coef_ = np.array([1.0, 2.0, 3.0])
+        mock_sklearn_model.intercept_ = -5.5
+        mock_model.regr = mock_sklearn_model
+
+        result = training_repo._extract_model_equation(mock_model, sample_params)
+
+        assert result['intercept'] == -5.5
+        assert 'target = -5.500000 +' in result['equation_string']
+
+
+class TestInitDataPreprocessor:
+    """Tests for _init_data_preprocessor method."""
+
+    def test_init_preprocessor_basic(self, training_repo, sample_params):
+        """Test basic preprocessor initialization."""
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor.target_variable == 'target'
+        assert preprocessor.input_columns == ['var1', 'var2', 'var3']
+
+    def test_init_preprocessor_with_scaler(self, training_repo, sample_params):
+        """Test preprocessor initialization with scaler enabled."""
+        sample_params.use_scaler = True
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor.scaler_name == 'Standard Scaler'
+
+    def test_init_preprocessor_without_scaler(self, training_repo, sample_params):
+        """Test preprocessor initialization without scaler."""
+        sample_params.use_scaler = False
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor.scaler_name == 'None'
+
+    def test_init_preprocessor_with_ar(self, training_repo, sample_params):
+        """Test preprocessor initialization with autoregressive variable."""
+        sample_params.include_ar = True
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor.ar_var == 'target'
+
+    def test_init_preprocessor_without_ar(self, training_repo, sample_params):
+        """Test preprocessor initialization without autoregressive variable."""
+        sample_params.include_ar = False
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor.ar_var is None
+
+    def test_init_preprocessor_with_static_removal(self, training_repo, sample_params):
+        """Test preprocessor with static window removal enabled."""
+        sample_params.rem_static_win = True
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor.static_threshold == 1
+
+    def test_init_preprocessor_lag_configuration(self, training_repo, sample_params):
+        """Test preprocessor lag configuration."""
+        sample_params.lag_train = 5
+        sample_params.lag_val = 3
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        # Check that lag dictionaries are created correctly
+        for col in sample_params.variable_columns:
+            assert preprocessor.lag_train[col] == 5
+            assert preprocessor.lag_transform[col] == 3
+
+
+class TestAfterTrainCalculation:
+    """Tests for after_train_calculation method."""
+
+    @pytest.fixture
+    def mock_train_result(self, sample_params, sample_linear_model):
+        """Create a mock TrainModelResult."""
+        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 to return a simple array
+        sample_linear_model.predict = MagicMock(return_value=np.array([450.0]))
+
+        return TrainModelResult(
+            params=sample_params,
+            process_data=MagicMock(),
+            x_train=x_train,
+            x_test=x_test,
+            y_train=y_train,
+            y_test=y_test,
+            regr=sample_linear_model,
+            scaler_dict={},
+        )
+
+    def test_after_train_adds_predictions(self, training_repo, sample_params, mock_train_result):
+        """Test that predictions are added to result."""
+        result = training_repo.after_train_calculation(sample_params, mock_train_result)
+
+        assert result.y_pred is not None
+        assert len(result.y_pred) == len(result.y_test)
+        assert result.y_pred.name == 'target_pred'
+
+    def test_after_train_calculates_metrics(self, training_repo, sample_params, mock_train_result):
+        """Test that metrics are calculated."""
+        result = training_repo.after_train_calculation(sample_params, mock_train_result)
+
+        assert result.mse_val is not None
+        assert result.mae_val is not None
+        assert result.r2_val is not None
+        assert isinstance(result.mse_val, (int, float))
+        assert isinstance(result.mae_val, (int, float))
+        assert isinstance(result.r2_val, (int, float))
+
+    def test_after_train_extracts_equation(self, training_repo, sample_params, mock_train_result):
+        """Test that equation is extracted after training."""
+        result = training_repo.after_train_calculation(sample_params, mock_train_result)
+
+        assert result.equation is not None
+        assert 'target_variable' in result.equation
+        assert 'coefficients' in result.equation
+        assert 'intercept' in result.equation
+        assert 'equation_string' in result.equation
+        assert 'latex_equation' in result.equation
+        assert 'model_type' in result.equation
+
+    def test_after_train_sorts_data(self, training_repo, sample_params, mock_train_result):
+        """Test that data is sorted by index."""
+        # Shuffle indices
+        mock_train_result.x_train = mock_train_result.x_train.sample(frac=1)
+        mock_train_result.y_train = mock_train_result.y_train.sample(frac=1)
+
+        result = training_repo.after_train_calculation(sample_params, mock_train_result)
+
+        assert result.x_train.index.is_monotonic_increasing
+        assert result.y_train.index.is_monotonic_increasing
+        assert result.x_test.index.is_monotonic_increasing
+        assert result.y_test.index.is_monotonic_increasing
+
+    def test_after_train_logs_success(
+        self, training_repo, mock_logger, sample_params, mock_train_result
+    ):
+        """Test that success is logged."""
+        training_repo.after_train_calculation(sample_params, mock_train_result)
+
+        mock_logger.info.assert_called()
+        assert any(
+            'Model metrics calculated successfully' in str(call)
+            for call in mock_logger.info.call_args_list
+        )
+