1097 lines
38 KiB
Python
1097 lines
38 KiB
Python
"""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 (
|
|
ALLOWED_FRONTEND_DATE_FORMATS,
|
|
FRONTEND_DATE_FORMAT_TO_STRFTIME,
|
|
DataPreprocessor,
|
|
LinearRegressionModel,
|
|
_frontend_date_format_to_strftime,
|
|
validate_frontend_date_format,
|
|
)
|
|
|
|
|
|
class TestFrontendDateFormatToStrftime:
|
|
"""Tests for _frontend_date_format_to_strftime (models module)."""
|
|
|
|
def test_none_or_empty_returns_none(self):
|
|
"""None or empty string returns None."""
|
|
assert _frontend_date_format_to_strftime(None) is None
|
|
assert _frontend_date_format_to_strftime('') is None
|
|
|
|
def test_all_six_allowed_formats_convert_correctly(self):
|
|
"""All allowed frontend formats map to expected strftime."""
|
|
for frontend_fmt, strftime_fmt in FRONTEND_DATE_FORMAT_TO_STRFTIME.items():
|
|
assert _frontend_date_format_to_strftime(frontend_fmt) == strftime_fmt
|
|
|
|
def test_invalid_format_raises(self):
|
|
"""Invalid format raises ValueError with allowed list in message."""
|
|
with raises(ValueError, match='Invalid date_format'):
|
|
_frontend_date_format_to_strftime('yyyy-MM-dd')
|
|
with raises(ValueError, match='Allowed formats'):
|
|
_frontend_date_format_to_strftime('invalid')
|
|
|
|
|
|
class TestValidateFrontendDateFormat:
|
|
"""Tests for validate_frontend_date_format."""
|
|
|
|
def test_none_or_empty_does_not_raise(self):
|
|
"""None or empty string does not raise."""
|
|
validate_frontend_date_format(None)
|
|
validate_frontend_date_format('')
|
|
validate_frontend_date_format(' ')
|
|
|
|
def test_allowed_formats_do_not_raise(self):
|
|
"""All allowed formats pass validation."""
|
|
for fmt in ALLOWED_FRONTEND_DATE_FORMATS:
|
|
validate_frontend_date_format(fmt)
|
|
|
|
def test_invalid_format_raises(self):
|
|
"""Invalid format raises ValueError."""
|
|
with raises(ValueError, match='Invalid date_format'):
|
|
validate_frontend_date_format('yyyy-MM-dd')
|
|
with raises(ValueError, match='Invalid date_format'):
|
|
validate_frontend_date_format('custom-bad-format')
|
|
|
|
|
|
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 infers variable_columns when not set."""
|
|
model = LinearRegressionModel(target_variable='target')
|
|
|
|
data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})
|
|
|
|
# Model should infer variable_columns from data (all columns except target)
|
|
result = model.fit(data)
|
|
assert result is model
|
|
assert model.variable_columns == ['var1']
|
|
|
|
|
|
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) == 8 # Now includes RANGE_SELECTION step
|
|
|
|
|
|
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)
|
|
|
|
|
|
# ============================================================================
|
|
# Additional tests for coverage - LinearRegressionModel
|
|
# ============================================================================
|
|
|
|
|
|
def test_linear_regression_model_fit_without_target_variable():
|
|
"""Test fit raises error when target_variable is not set."""
|
|
model = LinearRegressionModel()
|
|
data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})
|
|
|
|
with raises(ValueError, match='target_variable must be set before fitting'):
|
|
model.fit(data)
|
|
|
|
|
|
def test_linear_regression_model_fit_with_missing_columns():
|
|
"""Test fit raises error when variable_columns are missing from data."""
|
|
model = LinearRegressionModel(
|
|
target_variable='target', variable_columns=['var1', 'var_missing']
|
|
)
|
|
data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})
|
|
|
|
with raises(ValueError, match='Columns not found in input data'):
|
|
model.fit(data)
|
|
|
|
|
|
def test_linear_regression_model_fit_with_missing_target():
|
|
"""Test fit raises error when target_variable is not in data."""
|
|
model = LinearRegressionModel(target_variable='missing_target', variable_columns=['var1'])
|
|
data = pd.DataFrame({'var1': [1, 2, 3], 'other': [3, 5, 7]})
|
|
|
|
with raises(ValueError, match='Target variable missing_target not found in input data'):
|
|
model.fit(data)
|
|
|
|
|
|
def test_linear_regression_model_fit_with_inf_values():
|
|
"""Test fit handles infinite values by converting to NaN."""
|
|
model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2'])
|
|
# Include some inf values that will be converted to NaN and rows dropped
|
|
data = pd.DataFrame(
|
|
{
|
|
'var1': [1.0, 2.0, 3.0, np.inf, 5.0],
|
|
'var2': [2.0, 3.0, 4.0, 5.0, 6.0],
|
|
'target': [3.0, 5.0, 7.0, 9.0, 11.0],
|
|
}
|
|
)
|
|
|
|
model.fit(data)
|
|
|
|
# Model should fit successfully after removing row with inf
|
|
assert model.weights is not None
|
|
assert 'var1' in model.variable_columns
|
|
assert 'var2' in model.variable_columns
|
|
|
|
|
|
def test_linear_regression_model_fit_polynomial():
|
|
"""Test fit with polynomial features."""
|
|
model = LinearRegressionModel(target_variable='target', variable_columns=['var1'], degree=2)
|
|
data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]})
|
|
|
|
model.fit(data)
|
|
|
|
assert model.poly_feature_names is not None
|
|
assert len(model.poly_feature_names) > 1
|
|
|
|
|
|
def test_linear_regression_model_predict_polynomial():
|
|
"""Test predict with polynomial features."""
|
|
model = LinearRegressionModel(target_variable='target', variable_columns=['var1'], degree=2)
|
|
train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 4, 9, 16, 25]})
|
|
model.fit(train_data)
|
|
|
|
test_data = pd.DataFrame({'var1': [6, 7]})
|
|
predictions = model.predict(test_data)
|
|
|
|
assert isinstance(predictions, np.ndarray)
|
|
assert len(predictions) == 2
|
|
|
|
|
|
def test_linear_regression_model_create_poly_features_degree_1():
|
|
"""Test create_poly_features returns input unchanged when degree <= 1."""
|
|
model = LinearRegressionModel(degree=1)
|
|
data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})
|
|
|
|
result = model.create_poly_features(data, fit=True)
|
|
|
|
pd.testing.assert_frame_equal(result, data)
|
|
|
|
|
|
def test_linear_regression_model_create_poly_features_not_fitted():
|
|
"""Test create_poly_features raises error when not fitted and fit=False."""
|
|
model = LinearRegressionModel(degree=2)
|
|
data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})
|
|
|
|
with raises(ValueError, match='PolynomialFeatures not fitted'):
|
|
model.create_poly_features(data, fit=False)
|
|
|
|
|
|
def test_linear_regression_model_create_poly_features_transform():
|
|
"""Test create_poly_features with fit=False after fitting."""
|
|
model = LinearRegressionModel(degree=2)
|
|
train_data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})
|
|
model.create_poly_features(train_data, fit=True)
|
|
|
|
test_data = pd.DataFrame({'var1': [4, 5], 'var2': [7, 8]})
|
|
result = model.create_poly_features(test_data, fit=False)
|
|
|
|
assert isinstance(result, pd.DataFrame)
|
|
assert len(result.columns) > 2
|
|
|
|
|
|
def test_linear_regression_model_get_regressor():
|
|
"""Test get_regressor returns the underlying LinearRegression model."""
|
|
model = LinearRegressionModel()
|
|
|
|
regressor = model.get_regressor()
|
|
|
|
from sklearn.linear_model import LinearRegression
|
|
|
|
assert isinstance(regressor, LinearRegression)
|
|
|
|
|
|
# ============================================================================
|
|
# Additional tests for coverage - DataPreprocessor
|
|
# ============================================================================
|
|
|
|
|
|
@patch('model_manager.sientia.models.treat_nan')
|
|
def test_data_preprocessor_treat_discontinuities_linear_interpolation(mock_treat_nan):
|
|
"""Test treat_discontinuities with 'linear interpolation' treatment."""
|
|
preprocessor = DataPreprocessor(nan_treatment='linear interpolation')
|
|
data = pd.DataFrame({'col1': [1, np.nan, 3]})
|
|
expected_data = pd.DataFrame({'col1': [1.0, 2.0, 3.0]})
|
|
mock_treat_nan.return_value = expected_data
|
|
|
|
result = preprocessor.treat_discontinuities(data)
|
|
|
|
# Should map 'linear interpolation' to 'fill linear'
|
|
mock_treat_nan.assert_called_once_with(data, 'fill linear')
|
|
pd.testing.assert_frame_equal(result, expected_data)
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_start_date():
|
|
"""Test range_selection filters by start_date."""
|
|
preprocessor = DataPreprocessor(start_date='2023-01-02')
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3]},
|
|
index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
assert len(result) == 2
|
|
assert result.index[0] == pd.Timestamp('2023-01-02')
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_end_date():
|
|
"""Test range_selection filters by end_date."""
|
|
preprocessor = DataPreprocessor(end_date='2023-01-02')
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3]},
|
|
index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
assert len(result) == 2
|
|
assert result.index[-1] == pd.Timestamp('2023-01-02')
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_invalid_start_date():
|
|
"""Test range_selection handles invalid start_date gracefully."""
|
|
preprocessor = DataPreprocessor(start_date='invalid-date')
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3]},
|
|
index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
# Should skip filtering and return original data
|
|
assert len(result) == 3
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_invalid_end_date():
|
|
"""Test range_selection handles invalid end_date gracefully."""
|
|
preprocessor = DataPreprocessor(end_date='invalid-date')
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3]},
|
|
index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
# Should skip filtering and return original data
|
|
assert len(result) == 3
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_removed_intervals():
|
|
"""Test range_selection removes specified intervals."""
|
|
preprocessor = DataPreprocessor(removed_intervals=[['2023-01-02', '2023-01-03']])
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3, 4, 5]},
|
|
index=pd.to_datetime(
|
|
['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04', '2023-01-05']
|
|
),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
assert len(result) == 3
|
|
assert pd.Timestamp('2023-01-02') not in result.index
|
|
assert pd.Timestamp('2023-01-03') not in result.index
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_invalid_interval():
|
|
"""Test range_selection handles invalid interval dates gracefully."""
|
|
preprocessor = DataPreprocessor(removed_intervals=[['invalid', 'dates']])
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3]},
|
|
index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
# Should skip invalid interval and return original data
|
|
assert len(result) == 3
|
|
|
|
|
|
def test_data_preprocessor_range_selection_with_short_interval():
|
|
"""Test range_selection skips intervals with less than 2 elements."""
|
|
preprocessor = DataPreprocessor(removed_intervals=[['2023-01-02']])
|
|
data = pd.DataFrame(
|
|
{'col1': [1, 2, 3]},
|
|
index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
|
|
)
|
|
|
|
result = preprocessor.range_selection(data)
|
|
|
|
# Should skip short interval and return original data
|
|
assert len(result) == 3
|
|
|
|
|
|
def test_data_preprocessor_predict():
|
|
"""Test predict method removes target and preserves feature order."""
|
|
preprocessor = DataPreprocessor(
|
|
target_variable='target',
|
|
input_columns=['var1', 'var2'],
|
|
steps_order=['Discontinuity Treatment'],
|
|
)
|
|
train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
|
|
preprocessor.fit(train_x)
|
|
|
|
test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'target': [9, 11]})
|
|
result = preprocessor.predict(test_x)
|
|
|
|
assert 'target' not in result.columns
|
|
assert 'var1' in result.columns
|
|
assert 'var2' in result.columns
|
|
|
|
|
|
def test_data_preprocessor_predict_without_target():
|
|
"""Test predict when target is not in transformed data."""
|
|
preprocessor = DataPreprocessor(
|
|
target_variable='target',
|
|
input_columns=['var1', 'var2'],
|
|
steps_order=['Discontinuity Treatment'],
|
|
)
|
|
train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
|
|
preprocessor.fit(train_x)
|
|
|
|
# Include target in test data so transform works, predict will remove it
|
|
test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'target': [9, 11]})
|
|
result = preprocessor.predict(test_x)
|
|
|
|
# Target should be removed by predict
|
|
assert 'target' not in result.columns
|
|
assert 'var1' in result.columns
|
|
assert 'var2' in result.columns
|
|
|
|
|
|
def test_data_preprocessor_predict_preserves_feature_order():
|
|
"""Test predict preserves feature order from fit."""
|
|
preprocessor = DataPreprocessor(
|
|
target_variable='target',
|
|
input_columns=['var1', 'var2'],
|
|
steps_order=['Discontinuity Treatment'],
|
|
)
|
|
train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
|
|
preprocessor.fit(train_x)
|
|
|
|
# Test data has columns in different order
|
|
test_x = pd.DataFrame({'var2': [5, 6], 'var1': [4, 5], 'target': [9, 11]})
|
|
result = preprocessor.predict(test_x)
|
|
|
|
# Should have columns in same order as during fit
|
|
assert list(result.columns) == ['var1', 'var2']
|
|
|
|
|
|
# ============================================================================
|
|
# Additional tests for 100% coverage
|
|
# ============================================================================
|
|
|
|
|
|
def test_data_preprocessor_predict_target_not_in_columns():
|
|
"""Test predict when target_variable is not in transformed data columns."""
|
|
preprocessor = DataPreprocessor(
|
|
target_variable='target',
|
|
input_columns=['var1', 'var2'],
|
|
steps_order=['Discontinuity Treatment'],
|
|
)
|
|
train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
|
|
preprocessor.fit(train_x)
|
|
|
|
# Test data without target column - predict should still work
|
|
test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6]})
|
|
# Need to add target for transform to work, then it gets removed
|
|
test_x['target'] = [9, 11]
|
|
|
|
# Manually remove target before calling predict to test the branch
|
|
preprocessor_copy = DataPreprocessor(
|
|
target_variable='nonexistent_target',
|
|
input_columns=['var1', 'var2'],
|
|
steps_order=['Discontinuity Treatment'],
|
|
)
|
|
preprocessor_copy.fit(train_x.rename(columns={'target': 'nonexistent_target'}))
|
|
|
|
test_x_no_target = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'nonexistent_target': [9, 11]})
|
|
result = preprocessor_copy.predict(test_x_no_target)
|
|
|
|
assert 'var1' in result.columns
|
|
assert 'var2' in result.columns
|
|
|
|
|
|
def test_data_preprocessor_predict_without_fitted_feature_order():
|
|
"""Test predict when _fitted_feature_order is None."""
|
|
preprocessor = DataPreprocessor(
|
|
target_variable='target',
|
|
input_columns=['var1', 'var2'],
|
|
steps_order=['Discontinuity Treatment'],
|
|
)
|
|
train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
|
|
preprocessor.fit(train_x)
|
|
|
|
# Manually set _fitted_feature_order to None to test the branch
|
|
preprocessor._fitted_feature_order = None
|
|
|
|
test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'target': [9, 11]})
|
|
result = preprocessor.predict(test_x)
|
|
|
|
# Should still work, just without reordering
|
|
assert 'target' not in result.columns
|
|
assert 'var1' in result.columns
|
|
assert 'var2' in result.columns
|