"""Unit tests for TrainingRepository.""" from io import BytesIO from unittest.mock import MagicMock, 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, _apply_support_filters, _ensure_date_column_parsed, _frontend_format_to_strftime, ) @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', target_variable='target', variable_columns=['var1', 'var2', 'var3'], lag_train={'var1': 0, 'var2': 0, 'var3': 0}, lag_val={'var1': 0, 'var2': 0, 'var3': 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=[], model_name='Linear Regression', degree=1, interaction_only=False, nan_treatment='drop', start_date=None, end_date=None, scaler_name='None', support_filters={}, static_threshold=None, date_column=None, date_format=None, ) @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', target_variable='y', variable_columns=['x'], lag_train={'x': 0}, lag_val={'x': 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=[], model_name='Linear Regression', degree=1, interaction_only=False, nan_treatment='drop', start_date=None, end_date=None, scaler_name='None', support_filters={}, static_threshold=None, date_column=None, date_format=None, ) # 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 sample_params.scaler_name = 'Standard Scaler' 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 sample_params.scaler_name = 'None' 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 and no static_threshold.""" sample_params.rem_static_win = True sample_params.static_threshold = None preprocessor = training_repo._init_data_preprocessor(sample_params) assert preprocessor.static_threshold == 1 def test_init_preprocessor_with_static_removal_custom_threshold( self, training_repo, sample_params ): """Test preprocessor with static window removal and custom static_threshold.""" sample_params.rem_static_win = True sample_params.static_threshold = 500 preprocessor = training_repo._init_data_preprocessor(sample_params) assert preprocessor.static_threshold == 500 def test_init_preprocessor_without_static_removal_ignores_threshold( self, training_repo, sample_params ): """Test preprocessor without static removal ignores static_threshold.""" sample_params.rem_static_win = False sample_params.static_threshold = 500 preprocessor = training_repo._init_data_preprocessor(sample_params) assert preprocessor.static_threshold is None def test_init_preprocessor_lag_configuration(self, training_repo, sample_params): """Test preprocessor lag configuration.""" sample_params.lag_train = {'var1': 5, 'var2': 5, 'var3': 5} sample_params.lag_val = {'var1': 3, 'var2': 3, 'var3': 3} preprocessor = training_repo._init_data_preprocessor(sample_params) # Check that lag dictionaries are passed correctly assert preprocessor.lag_train == {'var1': 5, 'var2': 5, 'var3': 5} assert preprocessor.lag_transform == {'var1': 3, 'var2': 3, 'var3': 3} class TestGetStaticThreshold: """Tests for _get_static_threshold method.""" def test_get_static_threshold_rem_static_win_false(self, training_repo, sample_params): """Test returns None when rem_static_win is False.""" sample_params.rem_static_win = False sample_params.static_threshold = 500 result = training_repo._get_static_threshold(sample_params) assert result is None def test_get_static_threshold_rem_static_win_true_with_value( self, training_repo, sample_params ): """Test returns static_threshold value when rem_static_win is True and value is set.""" sample_params.rem_static_win = True sample_params.static_threshold = 500 result = training_repo._get_static_threshold(sample_params) assert result == 500 def test_get_static_threshold_rem_static_win_true_with_none(self, training_repo, sample_params): """Test returns 1 when rem_static_win is True and static_threshold is None.""" sample_params.rem_static_win = True sample_params.static_threshold = None result = training_repo._get_static_threshold(sample_params) assert result == 1 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 sientia_do.operations.normalization import MinMaxScaler from model_manager.sientia.models import DataPreprocessor 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 sientia_do.operations.normalization import Z_Scaler from model_manager.sientia.models import DataPreprocessor 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.""" @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, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4]) y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') y_test = pd.Series([400, 500], index=[3, 4], name='target') # Mock predict to return arrays with correct length based on input def mock_predict(data): if len(data) == 3: # x_train return np.array([150.0, 250.0, 350.0]) else: # x_test return np.array([450.0, 550.0]) sample_linear_model.predict = MagicMock(side_effect=mock_predict) 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 ) 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, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4]) y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') y_test = pd.Series([400, 500], index=[3, 4], name='target') # Mock predict to return arrays with correct length based on input def mock_predict(data): if len(data) == 3: # x_train return np.array([150.0, 250.0, 350.0]) else: # x_test return np.array([450.0, 550.0]) sample_linear_model.predict = MagicMock(side_effect=mock_predict) # Create mock scaler with denormalize methods mock_scaler = MagicMock() mock_scaler.denormalize_single_input = MagicMock(side_effect=lambda x, col: x * 2) # denormalize_predictions needs to return correct length based on input def mock_denormalize_predictions(arr, col): return arr * 2 mock_scaler.denormalize_predictions = MagicMock(side_effect=mock_denormalize_predictions) # 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, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4]) y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') y_test = pd.Series([400, 500], index=[3, 4], name='target') # Mock predict to return arrays with correct length based on input def mock_predict(data): if len(data) == 3: # x_train return np.array([150.0, 250.0, 350.0]) else: # x_test return np.array([450.0, 550.0]) sample_linear_model.predict = MagicMock(side_effect=mock_predict) # 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 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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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) 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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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) 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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) @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 # ============================================================================ # Tests for _configure_datetime_index # ============================================================================ class TestConfigureDatetimeIndex: """Tests for _configure_datetime_index method.""" @pytest.fixture def training_repo(self, mock_logger): """Create a TrainingRepository instance.""" return TrainingRepository(logger=mock_logger) @pytest.fixture def datetime_params(self): """Minimal params for _configure_datetime_index (date_column/date_format can be None).""" return TrainModelParams( experiment_run_id=1, experiment_name='test', target_variable='y', variable_columns=['x'], lag_train={'x': 0}, lag_val={'x': 0}, rem_static_win=False, low_lim={}, upp_lim={}, window=0, use_scaler=False, include_ar=False, train_size=80, shuffle=True, bucket_name='b', file_name='f.csv', line_separator=',', decimal_separator='.', removed_intervals=[], model_name='Linear Regression', degree=1, interaction_only=False, nan_treatment='drop', start_date=None, end_date=None, scaler_name='None', support_filters={}, static_threshold=None, date_column=None, date_format=None, ) def test_configure_datetime_index_already_datetime(self, training_repo, datetime_params): """Test _configure_datetime_index when index is already DatetimeIndex.""" data = pd.DataFrame( {'var1': [1, 2, 3], 'var2': [4, 5, 6]}, index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']), ) result = training_repo._configure_datetime_index(data, datetime_params) assert isinstance(result.index, pd.DatetimeIndex) assert len(result) == 3 def test_configure_datetime_index_with_timestamp_column(self, training_repo, datetime_params): """Test _configure_datetime_index with 'timestamp' column.""" data = pd.DataFrame( { 'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03'], 'var1': [1, 2, 3], 'var2': [4, 5, 6], } ) result = training_repo._configure_datetime_index(data, datetime_params) assert isinstance(result.index, pd.DatetimeIndex) assert 'timestamp' not in result.columns def test_configure_datetime_index_with_date_column(self, training_repo, datetime_params): """Test _configure_datetime_index with 'date' column.""" data = pd.DataFrame( { 'date': ['2023-01-01', '2023-01-02', '2023-01-03'], 'var1': [1, 2, 3], } ) result = training_repo._configure_datetime_index(data, datetime_params) assert isinstance(result.index, pd.DatetimeIndex) assert 'date' not in result.columns def test_configure_datetime_index_with_datetime_column(self, training_repo, datetime_params): """Test _configure_datetime_index with 'datetime' column.""" data = pd.DataFrame( { 'datetime': ['2023-01-01', '2023-01-02', '2023-01-03'], 'var1': [1, 2, 3], } ) result = training_repo._configure_datetime_index(data, datetime_params) assert isinstance(result.index, pd.DatetimeIndex) assert 'datetime' not in result.columns def test_configure_datetime_index_first_column_datetime(self, training_repo, datetime_params): """Test _configure_datetime_index when first column looks like datetime.""" data = pd.DataFrame( { 'my_date': ['2023-01-01', '2023-01-02', '2023-01-03'], 'var1': [1, 2, 3], } ) result = training_repo._configure_datetime_index(data, datetime_params) assert isinstance(result.index, pd.DatetimeIndex) assert 'my_date' not in result.columns @pytest.mark.filterwarnings('ignore::UserWarning') def test_configure_datetime_index_no_timestamp_column(self, training_repo, datetime_params): """Test _configure_datetime_index when no timestamp column found.""" data = pd.DataFrame( { 'var1': ['text_a', 'text_b', 'text_c'], 'var2': ['text_d', 'text_e', 'text_f'], } ) result = training_repo._configure_datetime_index(data, datetime_params) # Should return original data unchanged (no valid datetime columns) assert 'var1' in result.columns assert 'var2' in result.columns @pytest.mark.filterwarnings('ignore::UserWarning') def test_configure_datetime_index_invalid_timestamp_column( self, training_repo, datetime_params ): """Test _configure_datetime_index with invalid timestamp values.""" data = pd.DataFrame( { 'timestamp': ['not_a_date', 'also_not', 'nope'], 'var1': [1, 2, 3], } ) result = training_repo._configure_datetime_index(data, datetime_params) # Should skip invalid column and try first column assert 'var1' in result.columns @pytest.mark.filterwarnings('ignore::UserWarning') def test_configure_datetime_index_invalid_first_column(self, training_repo, datetime_params): """Test _configure_datetime_index when first column is not datetime.""" data = pd.DataFrame( { 'var1': ['a', 'b', 'c'], 'var2': [1, 2, 3], } ) result = training_repo._configure_datetime_index(data, datetime_params) # Should return original data unchanged assert 'var1' in result.columns assert 'var2' in result.columns def test_configure_datetime_index_first_column_all_nan(self, training_repo, datetime_params): """Test _configure_datetime_index when first column has all NaN values.""" data = pd.DataFrame( { 'first_col': [np.nan, np.nan, np.nan], 'var1': [1, 2, 3], } ) result = training_repo._configure_datetime_index(data, datetime_params) # Should return original data unchanged (first column has no valid values) assert 'first_col' in result.columns assert 'var1' in result.columns # ============================================================================ # Tests for _init_data_preprocessor with removed_intervals # ============================================================================ class TestInitDataPreprocessorWithRemovedIntervals: """Tests for _init_data_preprocessor with removed_intervals.""" @pytest.fixture def training_repo(self, mock_logger): """Create a TrainingRepository instance.""" return TrainingRepository(logger=mock_logger) def test_init_data_preprocessor_with_removed_intervals(self, training_repo, sample_params): """Test _init_data_preprocessor with removed_intervals.""" sample_params.removed_intervals = [ ['2023-01-01', '2023-01-02'], ['2023-02-01', '2023-02-02'], ] preprocessor = training_repo._init_data_preprocessor(sample_params) assert preprocessor is not None assert preprocessor.removed_intervals is not None assert len(preprocessor.removed_intervals) == 2 def test_init_data_preprocessor_with_tuple_intervals(self, training_repo, sample_params): """Test _init_data_preprocessor with tuple intervals.""" sample_params.removed_intervals = [('2023-01-01', '2023-01-02')] preprocessor = training_repo._init_data_preprocessor(sample_params) assert preprocessor is not None assert preprocessor.removed_intervals is not None # ============================================================================ # Tests for _extract_model_equation with polynomial features # ============================================================================ class TestExtractModelEquationPolynomial: """Tests for _extract_model_equation with polynomial features.""" @pytest.fixture def training_repo(self, mock_logger): """Create a TrainingRepository instance.""" return TrainingRepository(logger=mock_logger) def test_extract_model_equation_polynomial(self, training_repo, sample_params): """Test _extract_model_equation with polynomial features.""" # Create a mock regressor with polynomial features mock_regr = MagicMock() mock_regr.regr.coef_ = np.array([0.5, 0.3, 0.2]) mock_regr.regr.intercept_ = 1.0 mock_regr.poly_feature_names = ['var1', 'var2', 'var1^2'] sample_params.degree = 2 result = training_repo._extract_model_equation(mock_regr, sample_params) assert 'equation_string' in result assert 'latex_equation' in result assert 'var1' in result['equation_string'] def test_extract_model_equation_linear(self, training_repo, sample_params): """Test _extract_model_equation with linear features.""" mock_regr = MagicMock() mock_regr.regr.coef_ = np.array([0.5, 0.3]) mock_regr.regr.intercept_ = 1.0 mock_regr.poly_feature_names = None sample_params.degree = 1 result = training_repo._extract_model_equation(mock_regr, sample_params) assert 'equation_string' in result assert 'latex_equation' in result # ============================================================================ # Tests for _frontend_format_to_strftime and _ensure_date_column_parsed # ============================================================================ class TestFrontendFormatToStrftime: """Tests for _frontend_format_to_strftime (alias from models).""" def test_empty_or_none_returns_none(self): """Empty string or None returns None (same behavior as models module).""" assert _frontend_format_to_strftime('') is None assert _frontend_format_to_strftime(None) is None def test_allowed_formats_convert_correctly(self): """Allowed frontend formats convert to strftime via mapper.""" assert _frontend_format_to_strftime('dd/MM/yyyy HH:mm:ss') == '%d/%m/%Y %H:%M:%S' assert _frontend_format_to_strftime('yyyy-MM-dd HH:mm:ss') == '%Y-%m-%d %H:%M:%S' assert _frontend_format_to_strftime('MM/dd/yyyy HH:mm:ss') == '%m/%d/%Y %H:%M:%S' def test_invalid_format_raises(self): """Invalid date format raises ValueError.""" with pytest.raises(ValueError, match='Invalid date_format'): _frontend_format_to_strftime('yyyy-MM-dd') class TestEnsureDateColumnParsed: """Tests for _ensure_date_column_parsed.""" @pytest.fixture def date_params(self): """Params with date_column and date_format set.""" return TrainModelParams( experiment_run_id=1, experiment_name='test', target_variable='y', variable_columns=['x'], lag_train={'x': 0}, lag_val={'x': 0}, rem_static_win=False, low_lim={}, upp_lim={}, window=0, use_scaler=False, include_ar=False, train_size=80, shuffle=True, bucket_name='b', file_name='f.csv', line_separator=',', decimal_separator='.', removed_intervals=[], model_name='Linear Regression', degree=1, interaction_only=False, nan_treatment='drop', start_date=None, end_date=None, scaler_name='None', support_filters={}, static_threshold=None, date_column='ts', date_format='yyyy-MM-dd HH:mm:ss', ) def test_returns_unchanged_when_no_date_column(self, date_params): """When params.date_column is None, data is returned unchanged.""" date_params.date_column = None date_params.date_format = None data = pd.DataFrame({'ts': ['2023-01-01'], 'x': [1]}) result = _ensure_date_column_parsed(data, date_params) pd.testing.assert_frame_equal(result, data) def test_returns_unchanged_when_column_missing(self, date_params): """When date_column not in data columns, data is returned unchanged.""" data = pd.DataFrame({'other': [1], 'x': [2]}) result = _ensure_date_column_parsed(data, date_params) pd.testing.assert_frame_equal(result, data) def test_parses_column_with_format(self, date_params): """When date_column and date_format set, column is parsed as datetime.""" data = pd.DataFrame( { 'ts': ['2023-01-01 10:00:00', '2023-06-15 14:30:00'], 'x': [1, 2], } ) result = _ensure_date_column_parsed(data, date_params) assert result['ts'].dtype == 'datetime64[ns]' assert result['ts'].iloc[0].year == 2023 assert result['ts'].iloc[0].month == 1 assert result['ts'].iloc[1].month == 6 def test_invalid_values_coerced_to_nat(self, date_params): """Invalid date strings are coerced to NaT when format is set.""" date_params.date_format = 'yyyy-MM-dd HH:mm:ss' data = pd.DataFrame( { 'ts': ['2023-01-01 00:00:00', 'not-a-date', '2023-12-31 00:00:00'], 'x': [1, 2, 3], } ) result = _ensure_date_column_parsed(data, date_params) assert pd.isna(result['ts'].iloc[1]) assert result['ts'].iloc[0].year == 2023 assert result['ts'].iloc[2].month == 12 # ============================================================================ # Tests for _apply_support_filters # ============================================================================ class TestApplySupportFilters: """Tests for _apply_support_filters function.""" def test_empty_support_filters_returns_unchanged(self): """When support_filters is empty, data_view is returned unchanged.""" data = pd.DataFrame({'x': [1, 2, 3], 'target': [10, 20, 30]}) result = _apply_support_filters(data, 'target', {}) pd.testing.assert_frame_equal(result, data) def test_target_not_in_columns_returns_unchanged(self): """When target_variable is not in data_view columns, return unchanged.""" data = pd.DataFrame({'x': [1, 2, 3], 'y': [10, 20, 30]}) result = _apply_support_filters(data, 'target', {'x': {}}) pd.testing.assert_frame_equal(result, data) def test_variable_not_in_columns_skipped(self): """When a filter variable is not in data_view, that variable is skipped.""" data = pd.DataFrame({'x': [1, 2, 3], 'target': [10, 20, 30]}) support_filters = { 'missing_var': { 'upper_line': {'intercept': 100, 'angle': 10}, 'lower_line': {'intercept': 0, 'angle': -10}, }, } result = _apply_support_filters(data, 'target', support_filters) pd.testing.assert_frame_equal(result, data) def test_snake_case_upper_lower_line(self): """Support filters with upper_line/lower_line (snake_case) filter rows.""" data = pd.DataFrame( { 'x': [1.0, 2.0, 3.0, 4.0], 'target': [2.0, 4.0, 6.0, 8.0], } ) support_filters = { 'x': { 'upper_line': {'intercept': 1.0, 'angle': 50}, 'lower_line': {'intercept': -1.0, 'angle': -50}, }, } result = _apply_support_filters(data, 'target', support_filters) assert len(result) <= 4 assert list(result.columns) == ['x', 'target'] def test_camel_case_upper_lower_line(self): """Support filters with upperLine/lowerLine (camelCase) are accepted.""" data = pd.DataFrame( { 'x': [1.0, 2.0, 3.0], 'target': [1.0, 2.0, 3.0], } ) support_filters = { 'x': { 'upperLine': {'intercept': 2, 'angle': 5}, 'lowerLine': {'intercept': 0, 'angle': -5}, }, } result = _apply_support_filters(data, 'target', support_filters) assert len(result) <= 3 assert list(result.columns) == ['x', 'target'] def test_two_variables_ands_masks(self): """Two variables apply AND of both masks.""" data = pd.DataFrame( { 'a': [1.0, 2.0, 3.0], 'b': [1.0, 2.0, 3.0], 'target': [2.0, 2.0, 2.0], } ) support_filters = { 'a': { 'upper_line': {'intercept': 10, 'angle': 45}, 'lower_line': {'intercept': -10, 'angle': -45}, }, 'b': { 'upper_line': {'intercept': 10, 'angle': 45}, 'lower_line': {'intercept': -10, 'angle': -45}, }, } result = _apply_support_filters(data, 'target', support_filters) assert len(result) <= 3 assert list(result.columns) == ['a', 'b', 'target'] def test_missing_upper_or_lower_skips_variable(self): """If upper_line or lower_line is missing, that variable is skipped.""" data = pd.DataFrame({'x': [1, 2, 3], 'target': [10, 20, 30]}) support_filters = { 'x': {'upper_line': {'intercept': 100, 'angle': 0}}, } result = _apply_support_filters(data, 'target', support_filters) pd.testing.assert_frame_equal(result, data)