06571011f2
This parameter allows customizing the threshold (1-1000) used when rem_static_win is enabled, defaulting to 1 if null. Updates include parameter definition, business rule validation, repository logic for passing the threshold, documentation in README.md and PIPELINE_PARAMS_CHANGELOG.md, and new unit and integration tests.
1042 lines
40 KiB
Python
1042 lines
40 KiB
Python
"""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
|
|
|
|
|
|
@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,
|
|
)
|
|
|
|
|
|
@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,
|
|
)
|
|
|
|
# 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 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: 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: 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: 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: 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: 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: 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: 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: 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)
|
|
|
|
def test_configure_datetime_index_already_datetime(self, training_repo):
|
|
"""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)
|
|
|
|
assert isinstance(result.index, pd.DatetimeIndex)
|
|
assert len(result) == 3
|
|
|
|
def test_configure_datetime_index_with_timestamp_column(self, training_repo):
|
|
"""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)
|
|
|
|
assert isinstance(result.index, pd.DatetimeIndex)
|
|
assert 'timestamp' not in result.columns
|
|
|
|
def test_configure_datetime_index_with_date_column(self, training_repo):
|
|
"""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)
|
|
|
|
assert isinstance(result.index, pd.DatetimeIndex)
|
|
assert 'date' not in result.columns
|
|
|
|
def test_configure_datetime_index_with_datetime_column(self, training_repo):
|
|
"""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)
|
|
|
|
assert isinstance(result.index, pd.DatetimeIndex)
|
|
assert 'datetime' not in result.columns
|
|
|
|
def test_configure_datetime_index_first_column_datetime(self, training_repo):
|
|
"""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)
|
|
|
|
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):
|
|
"""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)
|
|
|
|
# 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):
|
|
"""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)
|
|
|
|
# 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):
|
|
"""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)
|
|
|
|
# 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):
|
|
"""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)
|
|
|
|
# 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
|