"""Unit tests for sientia models module.""" from unittest.mock import MagicMock, patch import numpy as np import pandas as pd from pytest import raises from model_manager.sientia.models import ( DataPreprocessor, LinearRegressionModel, ) class _IterableWithContains: def __init__(self, iterable, contains_values): self._iterable = iterable self._contains = set(contains_values) def __iter__(self): return iter(self._iterable) def __contains__(self, item): return item in self._contains # LinearRegressionModel Tests def test_linear_regression_model_init_default(): """Test LinearRegressionModel initialization with default parameters.""" model = LinearRegressionModel() assert model.target_variable == '' assert model.variable_columns is None assert model.model_params is None assert model.clipping is None assert model.weights is None assert model.q1_target is None assert model.q3_target is None def test_linear_regression_model_init_with_params(): """Test LinearRegressionModel initialization with parameters.""" target = 'target' variables = ['var1', 'var2'] params = {'fit_intercept': True} clipping = {'min': 0, 'max': 100} weights = {'var1': 0.5, 'var2': 0.3} model = LinearRegressionModel( target_variable=target, variable_columns=variables, model_params=params, clipping=clipping, weights=weights, ) assert model.target_variable == target assert model.variable_columns == variables assert model.model_params == params assert model.clipping == clipping assert model.weights == weights def test_linear_regression_model_fit(): """Test LinearRegressionModel fit method.""" model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2']) data = pd.DataFrame( {'var1': [1, 2, 3, 4, 5], 'var2': [2, 3, 4, 5, 6], 'target': [3, 5, 7, 9, 11]} ) result = model.fit(data) assert result is model assert model.q1_target is not None assert model.q3_target is not None assert model.weights is not None assert 'Bias' in model.weights def test_linear_regression_model_fit_without_variable_columns(): """Test LinearRegressionModel fit raises AssertionError without variable_columns.""" model = LinearRegressionModel(target_variable='target') data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) with raises(AssertionError, match='variable_columns must be set before fitting'): model.fit(data) def test_linear_regression_model_predict_without_clipping(): """Test LinearRegressionModel predict without clipping.""" model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2']) train_data = pd.DataFrame( {'var1': [1, 2, 3, 4, 5], 'var2': [2, 3, 4, 5, 6], 'target': [3, 5, 7, 9, 11]} ) model.fit(train_data) test_data = pd.DataFrame({'var1': [6, 7], 'var2': [7, 8]}) predictions = model.predict(test_data) assert isinstance(predictions, np.ndarray) assert len(predictions) == 2 def test_linear_regression_model_predict_with_clipping_max(): """Test LinearRegressionModel predict with clipping max.""" model = LinearRegressionModel( target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 5} ) train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]}) model.fit(train_data) test_data = pd.DataFrame({'var1': [10]}) predictions = model.predict(test_data) assert predictions[0] == model.q3_target def test_linear_regression_model_predict_with_clipping_min(): """Test LinearRegressionModel predict with clipping min.""" model = LinearRegressionModel( target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 10} ) train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]}) model.fit(train_data) test_data = pd.DataFrame({'var1': [-10]}) predictions = model.predict(test_data) assert predictions[0] == model.q1_target def test_linear_regression_model_predict_with_clipping_within_range(): """Test LinearRegressionModel predict with clipping but value within range.""" model = LinearRegressionModel( target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 10} ) train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]}) model.fit(train_data) test_data = pd.DataFrame({'var1': [3]}) predictions = model.predict(test_data) # Prediction should be within range and not clipped assert 0 <= predictions[0] <= 10 # DataPreprocessor Tests def test_data_preprocessor_init_default(): """Test DataPreprocessor initialization with default parameters.""" preprocessor = DataPreprocessor() assert preprocessor.date_column == '' assert preprocessor.target_variable == '' assert preprocessor.input_columns is None assert preprocessor.nan_treatment is None assert preprocessor.lag_train == {} assert preprocessor.lag_transform == {} assert preprocessor.scaler is None def test_data_preprocessor_init_with_standard_scaler(): """Test DataPreprocessor initialization with Standard Scaler.""" preprocessor = DataPreprocessor(scaler_name='Standard Scaler') assert preprocessor.scaler is not None def test_data_preprocessor_init_with_none_scaler(): """Test DataPreprocessor initialization with None scaler.""" preprocessor = DataPreprocessor(scaler_name='None') assert preprocessor.scaler is None def test_data_preprocessor_init_with_unknown_scaler(): """Test DataPreprocessor initialization with unknown scaler.""" preprocessor = DataPreprocessor(scaler_name='Unknown') assert preprocessor.scaler is None def test_data_preprocessor_init_with_custom_steps_order(): """Test DataPreprocessor initialization with custom steps order.""" custom_steps = ['Normalization', 'Feature Creation'] preprocessor = DataPreprocessor(steps_order=custom_steps) assert 'Normalization' in preprocessor.steps_order assert 'Feature Creation' in preprocessor.steps_order assert len(preprocessor.steps_order) == 7 def test_data_preprocessor_get_scaler(): """Test DataPreprocessor get_scaler method.""" preprocessor = DataPreprocessor(scaler_name='Standard Scaler') scaler = preprocessor.get_scaler() assert scaler is not None @patch('model_manager.sientia.models.treat_nan') def test_data_preprocessor_treat_discontinuities_with_treatment(mock_treat_nan): """Test treat_discontinuities with nan_treatment.""" preprocessor = DataPreprocessor(nan_treatment='drop') data = pd.DataFrame({'col1': [1, 2, np.nan]}) expected_data = pd.DataFrame({'col1': [1, 2]}) mock_treat_nan.return_value = expected_data result = preprocessor.treat_discontinuities(data) mock_treat_nan.assert_called_once_with(data, 'drop') pd.testing.assert_frame_equal(result, expected_data) def test_data_preprocessor_treat_discontinuities_without_treatment(): """Test treat_discontinuities without nan_treatment.""" preprocessor = DataPreprocessor() data = pd.DataFrame({'col1': [1, 2, 3]}) result = preprocessor.treat_discontinuities(data) pd.testing.assert_frame_equal(result, data) def test_data_preprocessor_lag_selection_with_lag(): """Test lag_selection with lag.""" preprocessor = DataPreprocessor() data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) lag_dict = {'var1': 1} result = preprocessor.lag_selection(data, lag_dict) assert len(result) == 4 assert result['var1'].iloc[0] == 1 def test_data_preprocessor_lag_selection_without_lag(): """Test lag_selection without lag.""" preprocessor = DataPreprocessor() data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) lag_dict = {} result = preprocessor.lag_selection(data, lag_dict) assert len(result) == 5 def test_data_preprocessor_lag_selection_with_zero_lag(): """Test lag_selection with zero lag.""" preprocessor = DataPreprocessor() data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) lag_dict = {'var1': 0} result = preprocessor.lag_selection(data, lag_dict) assert len(result) == 5 @patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer') def test_data_preprocessor_treat_static_windows(mock_analyzer_class): """Test treat_static_windows method.""" preprocessor = DataPreprocessor(static_threshold=3) data = pd.DataFrame({'col1': [1, 1, 1, 2, 3]}) mock_analyzer = MagicMock() mock_analyzer_class.return_value = mock_analyzer mock_analyzer.get_treated_data.return_value = data preprocessor.treat_static_windows(data) mock_analyzer.infer_frequency.assert_called_once() assert mock_analyzer.identify_static_windows.called assert mock_analyzer.treat_static_windows.called def test_data_preprocessor_treat_static_windows_without_threshold(): """Test treat_static_windows without threshold.""" preprocessor = DataPreprocessor() data = pd.DataFrame({'col1': [1, 2, 3]}) result = preprocessor.treat_static_windows(data) pd.testing.assert_frame_equal(result, data) @patch('model_manager.sientia.models.limit_dataset') def test_data_preprocessor_adjust_limits(mock_limit_dataset): """Test adjust_limits method.""" preprocessor = DataPreprocessor(low_lim={'col1': 0}, upp_lim={'col1': 10}) data = pd.DataFrame({'col1': [1, 2, 3]}) expected_data = pd.DataFrame({'col1': [1, 2, 3]}) mock_limit_dataset.return_value = (expected_data, {'col1': 0}, {'col1': 10}) result = preprocessor.adjust_limits(data) mock_limit_dataset.assert_called_once() pd.testing.assert_frame_equal(result, expected_data) @patch('model_manager.sientia.models.create_features') def test_data_preprocessor_create_features(mock_create_features): """Test create_features method.""" preprocessor = DataPreprocessor( self_operations=['{var1}_{pow}_{2}'], cross_operations=['{var1}_{*}_{var2}'] ) data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4]}) expected_data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var1_pow_2': [1, 4, 9]}) mock_create_features.return_value = expected_data result = preprocessor.create_features(data) mock_create_features.assert_called_once() pd.testing.assert_frame_equal(result, expected_data) def test_data_preprocessor_create_ar(): """Test create_ar method.""" preprocessor = DataPreprocessor(target_variable='target', ar_var='ar_target') data = pd.DataFrame({'target': [1, 2, 3, 4, 5]}) result = preprocessor.create_ar(data) assert 'ar_target' in result.columns assert len(result) == 4 def test_data_preprocessor_create_ar_without_ar_var(): """Test create_ar without ar_var.""" preprocessor = DataPreprocessor(target_variable='target') data = pd.DataFrame({'target': [1, 2, 3, 4, 5]}) result = preprocessor.create_ar(data) assert len(result) == 5 def test_data_preprocessor_create_lags(): """Test create_lags method.""" preprocessor = DataPreprocessor(created_lags={'var1': 1}) data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) result = preprocessor.create_lags(data) assert 'var1_lag1' in result.columns assert len(result) == 4 def test_data_preprocessor_create_lags_with_zero_lag(): """Test create_lags with zero lag.""" preprocessor = DataPreprocessor(created_lags={'var1': 0}) data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) result = preprocessor.create_lags(data) assert 'var1_lag0' not in result.columns assert len(result) == 5 def test_data_preprocessor_create_lags_without_created_lags(): """Test create_lags without created_lags.""" preprocessor = DataPreprocessor() data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) result = preprocessor.create_lags(data) assert len(result) == 5 def test_data_preprocessor_create_lags_with_missing_column(): """Test create_lags with missing column.""" preprocessor = DataPreprocessor(created_lags={'var2': 1}) data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) result = preprocessor.create_lags(data) assert 'var2_lag1' not in result.columns def test_data_preprocessor_fit_with_x_and_y(): """Test fit method with x and y.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1', 'var2'], steps_order=['Discontinuity Treatment'], ) x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4]}) y = pd.Series([3, 5, 7], name='target') result = preprocessor.fit(x, y) assert result is preprocessor def test_data_preprocessor_fit_with_only_x(): """Test fit method with only x.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1', 'var2'], steps_order=['Discontinuity Treatment'], ) x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]}) result = preprocessor.fit(x) assert result is preprocessor def test_data_preprocessor_fit_without_data(): """Test fit method without data.""" preprocessor = DataPreprocessor(target_variable='target', input_columns=['var1']) with raises(ValueError, match='No data was provided'): preprocessor.fit(None, None) def test_data_preprocessor_fit_without_input_columns(): """Test fit method without input_columns.""" preprocessor = DataPreprocessor(target_variable='target') x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) with raises(AssertionError, match='input_columns must be set'): preprocessor.fit(x) def test_data_preprocessor_fit_with_normalization(): """Test fit method with normalization.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1', 'var2'], scaler_name='Standard Scaler', scaler_params={}, steps_order=['Normalization'], ) x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]}) result = preprocessor.fit(x) assert result is preprocessor assert preprocessor.scaler_params is not None def test_data_preprocessor_transform_with_timestamp(): """Test transform method with timestamp column.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1'], steps_order=['Discontinuity Treatment'] ) x = pd.DataFrame({'timestamp': [1, 2, 3], 'var1': [1, 2, 3], 'target': [3, 5, 7]}) result = preprocessor.transform(x) assert 'timestamp' not in result.columns def test_data_preprocessor_transform_without_timestamp(): """Test transform method without timestamp column.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1'], steps_order=['Discontinuity Treatment'] ) x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) result = preprocessor.transform(x) assert 'var1' in result.columns def test_data_preprocessor_transform_without_input_columns(): """Test transform method without input_columns.""" preprocessor = DataPreprocessor(target_variable='target') x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) with raises(AssertionError, match='input_columns must be set'): preprocessor.transform(x) def test_data_preprocessor_transform_with_feature_creation(): """Test transform method with feature creation.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1'], self_operations=['{var1}_{pow}_{2}'], steps_order=['Feature Creation'], ) x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) with patch('model_manager.sientia.models.create_features') as mock_create: mock_create.return_value = x preprocessor.transform(x) mock_create.assert_called_once() def test_data_preprocessor_transform_with_lag_creation(): """Test transform method with lag creation.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1', 'ar_target'], ar_var='ar_target', created_lags={'var1': 1}, steps_order=['Lag Creation'], ) x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]}) result = preprocessor.transform(x) assert 'ar_target' in result.columns assert 'var1_lag1' in result.columns def test_data_preprocessor_transform_with_normalization(): """Test transform method with normalization.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1'], scaler_name='Standard Scaler', scaler_params={}, steps_order=['Normalization'], ) train_x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) preprocessor.fit(train_x) test_x = pd.DataFrame({'var1': [4, 5, 6], 'target': [9, 11, 13]}) result = preprocessor.transform(test_x) assert 'var1' in result.columns def test_data_preprocessor_get_required_columns_with_self_operations(): """Test get_required_columns with self_operations.""" preprocessor = DataPreprocessor(self_operations=['{var1}_{pow}_{2}']) existing_columns = ['var1', 'var2'] result = preprocessor.get_required_columns(existing_columns) assert 'var1' in result def test_data_preprocessor_get_required_columns_with_cross_operations(): """Test get_required_columns with cross_operations.""" preprocessor = DataPreprocessor(cross_operations=['{var1}_{*}_{var2}']) existing_columns = ['var1', 'var2'] result = preprocessor.get_required_columns(existing_columns) assert 'var1' in result assert 'var2' in result def test_data_preprocessor_get_required_columns_with_created_lags(): """Test get_required_columns with created_lags.""" preprocessor = DataPreprocessor(created_lags={'var1': 1}) existing_columns = ['var1', 'var2'] result = preprocessor.get_required_columns(existing_columns) assert 'var1' in result def test_data_preprocessor_get_required_columns_with_missing_columns(): """Test get_required_columns with missing columns in existing_columns.""" preprocessor = DataPreprocessor( self_operations=['{var3}_{pow}_{2}'], cross_operations=['{var4}_{*}_{var5}'] ) existing_columns = ['var1', 'var2'] result = preprocessor.get_required_columns(existing_columns) assert 'var3' in result assert 'var4' in result assert 'var5' in result def test_data_preprocessor_get_required_columns_removes_duplicates(): """Test get_required_columns removes duplicates from self_operations.""" preprocessor = DataPreprocessor(self_operations=['{var1}_{pow}_{2}'], created_lags={'var1': 1}) existing_columns = ['var1', 'var2'] result = preprocessor.get_required_columns(existing_columns) # var1 is in existing_columns, so it should not be in required_columns assert 'var1' not in result or result.count('var1') <= 1 def test_data_preprocessor_get_required_columns_removes_self_operations_branch(): """Ensure line 278 removes columns present in self_operations iterable.""" preprocessor = DataPreprocessor( self_operations=_IterableWithContains(['{var1}_{pow}_{2}'], contains_values=['var1']) ) existing_columns: list[str] = [] result = preprocessor.get_required_columns(existing_columns) assert 'var1' not in result def test_data_preprocessor_get_required_columns_removes_cross_operations_branch(): """Ensure line 285 removes columns present in cross_operations iterable.""" preprocessor = DataPreprocessor( cross_operations=_IterableWithContains(['{var1}_{*}_{var2}'], contains_values=['var1']) ) existing_columns: list[str] = [] result = preprocessor.get_required_columns(existing_columns) assert 'var1' not in result assert 'var2' in result def test_data_preprocessor_get_required_columns_removes_created_lags(): """Test get_required_columns removes columns from created_lags when column is in created_lags dict - covers line 292.""" preprocessor = DataPreprocessor(created_lags={'var1': 1, 'var2': 1}) existing_columns = ['var3'] result = preprocessor.get_required_columns(existing_columns) # var1 and var2 should be removed because they're in created_lags dict and not in existing_columns assert 'var1' not in result assert 'var2' not in result def test_data_preprocessor_fit_all_steps(): """Test fit method with all steps.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1'], nan_treatment='drop', lag_train={'var1': 1}, static_threshold=3, low_lim={'var1': 0}, upp_lim={'var1': 10}, scaler_name='Standard Scaler', scaler_params={}, ) x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]}) with patch('model_manager.sientia.models.treat_nan') as mock_treat: with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'): with patch('model_manager.sientia.models.limit_dataset') as mock_limit: mock_treat.return_value = x mock_limit.return_value = (x, {'var1': 0}, {'var1': 10}) result = preprocessor.fit(x) assert result is preprocessor def test_data_preprocessor_transform_all_steps(): """Test transform method with all steps.""" preprocessor = DataPreprocessor( target_variable='target', input_columns=['var1'], nan_treatment='drop', lag_transform={'var1': 1}, static_threshold=3, low_lim={'var1': 0}, upp_lim={'var1': 10}, scaler_name='Standard Scaler', scaler_params={}, self_operations=['{var1}_{pow}_{2}'], ar_var='ar_target', created_lags={'var1': 1}, ) train_x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]}) with patch('model_manager.sientia.models.treat_nan') as mock_treat: with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'): with patch('model_manager.sientia.models.limit_dataset') as mock_limit: with patch('model_manager.sientia.models.create_features') as mock_create: mock_treat.return_value = train_x mock_limit.return_value = (train_x, {'var1': 0}, {'var1': 10}) mock_create.return_value = train_x preprocessor.fit(train_x) test_x = pd.DataFrame({'var1': [6, 7, 8, 9, 10], 'target': [13, 15, 17, 19, 21]}) with patch('model_manager.sientia.models.treat_nan') as mock_treat: with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'): with patch('model_manager.sientia.models.limit_dataset') as mock_limit: with patch('model_manager.sientia.models.create_features') as mock_create: mock_treat.return_value = test_x mock_limit.return_value = (test_x, {'var1': 0}, {'var1': 10}) mock_create.return_value = test_x result = preprocessor.transform(test_x) assert isinstance(result, pd.DataFrame)