diff --git a/tests/utils/repository/test_training_repository.py b/tests/utils/repository/test_training_repository.py new file mode 100644 index 0000000..cea9ace --- /dev/null +++ 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 + ) +