From df6bf1dabae33db65e2cc4fbead2b281206a1899 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Fri, 19 Dec 2025 16:19:10 -0300 Subject: [PATCH] SIENTIAPDE-1430: Refactor static threshold calculation logic and enhance test coverage --- .../utils/repository/model_repository.py | 12 ++-- .../utils/repository/training_repository.py | 18 +++++- .../utils/repository/test_model_repository.py | 64 +++++++++++++++++++ .../repository/test_training_repository.py | 33 ++++++++++ 4 files changed, 118 insertions(+), 9 deletions(-) diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py index 123c54b..dcd1ba8 100644 --- a/model_manager/utils/repository/model_repository.py +++ b/model_manager/utils/repository/model_repository.py @@ -205,12 +205,12 @@ class ModelRepository: self.model_serving.log_param('nan_treatment', data.params.nan_treatment) self.model_serving.log_param('lag_train', data.params.lag_train) self.model_serving.log_param('lag_transform', data.params.lag_val) - self.model_serving.log_param( - 'static_threshold', - (data.params.static_threshold if data.params.static_threshold is not None else 1) - if data.params.rem_static_win - else None, - ) + static_threshold_value = None + if data.params.rem_static_win: + static_threshold_value = ( + data.params.static_threshold if data.params.static_threshold is not None else 1 + ) + self.model_serving.log_param('static_threshold', static_threshold_value) self.model_serving.log_param('lower_limits', data.params.low_lim) self.model_serving.log_param('upper_limits', data.params.upp_lim) self.model_serving.log_param('scaler_name', data.params.scaler_name) diff --git a/model_manager/utils/repository/training_repository.py b/model_manager/utils/repository/training_repository.py index 2c6e396..7d5249c 100644 --- a/model_manager/utils/repository/training_repository.py +++ b/model_manager/utils/repository/training_repository.py @@ -242,6 +242,20 @@ class TrainingRepository: return scaler_dict + def _get_static_threshold(self, params: TrainModelParams) -> int | None: + """ + Get the static threshold value based on parameters. + + Args: + params: Training parameters containing static window configuration + + Returns: + int | None: Static threshold value (1-1000) if rem_static_win is True, None otherwise + """ + if not params.rem_static_win: + return None + return params.static_threshold if params.static_threshold is not None else 1 + def _init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor: """ Initialize DataPreprocessor with training parameters. @@ -269,9 +283,7 @@ class TrainingRepository: start_date=params.start_date, end_date=params.end_date, removed_intervals=removed_intervals, - static_threshold=(params.static_threshold if params.static_threshold is not None else 1) - if params.rem_static_win - else None, + static_threshold=self._get_static_threshold(params), low_lim=params.low_lim, upp_lim=params.upp_lim, scaler_name=params.scaler_name, diff --git a/tests/utils/repository/test_model_repository.py b/tests/utils/repository/test_model_repository.py index 00ab54a..165a86c 100644 --- a/tests/utils/repository/test_model_repository.py +++ b/tests/utils/repository/test_model_repository.py @@ -493,6 +493,70 @@ def test_save_run_with_equation( assert '/tmp/test_run/model_equation.json' in logged_artifacts # noqa: S108 +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.path.exists') +def test_save_run_with_static_threshold_value( + mock_exists, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _save_run logs static_threshold when rem_static_win is True and value is set.""" + from model_manager.utils.repository.model_repository import ModelRepository + + mock_model_serving_instance = MagicMock() + mock_model_serving_class.return_value = mock_model_serving_instance + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # Set static_threshold to a specific value + mock_train_result.params.rem_static_win = True + mock_train_result.params.static_threshold = 500 + mock_train_result.equation_path = None + + # Mock all path.exists calls to return True + mock_exists.return_value = True + + repo._save_run(mock_train_result) + + # Verify static_threshold was logged with the correct value + log_param_calls = { + call[0][0]: call[0][1] for call in mock_model_serving_instance.log_param.call_args_list + } + assert log_param_calls['static_threshold'] == 500 + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.path.exists') +def test_save_run_with_rem_static_win_false( + mock_exists, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _save_run logs static_threshold as None when rem_static_win is False.""" + from model_manager.utils.repository.model_repository import ModelRepository + + mock_model_serving_instance = MagicMock() + mock_model_serving_class.return_value = mock_model_serving_instance + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # Set rem_static_win to False + mock_train_result.params.rem_static_win = False + mock_train_result.params.static_threshold = 500 # Should be ignored + mock_train_result.equation_path = None + + # Mock all path.exists calls to return True + mock_exists.return_value = True + + repo._save_run(mock_train_result) + + # Verify static_threshold was logged as None + log_param_calls = { + call[0][0]: call[0][1] for call in mock_model_serving_instance.log_param.call_args_list + } + assert log_param_calls['static_threshold'] is None + + @patch('model_manager.utils.repository.model_repository.ModelServing') @patch('model_manager.utils.repository.model_repository.path.exists') def test_save_run_without_equation( diff --git a/tests/utils/repository/test_training_repository.py b/tests/utils/repository/test_training_repository.py index 443f61b..42fed41 100644 --- a/tests/utils/repository/test_training_repository.py +++ b/tests/utils/repository/test_training_repository.py @@ -261,6 +261,39 @@ class TestInitDataPreprocessor: assert preprocessor.lag_transform == {'var1': 3, 'var2': 3, 'var3': 3} +class TestGetStaticThreshold: + """Tests for _get_static_threshold method.""" + + def test_get_static_threshold_rem_static_win_false(self, training_repo, sample_params): + """Test returns None when rem_static_win is False.""" + sample_params.rem_static_win = False + sample_params.static_threshold = 500 + + result = training_repo._get_static_threshold(sample_params) + + assert result is None + + def test_get_static_threshold_rem_static_win_true_with_value( + self, training_repo, sample_params + ): + """Test returns static_threshold value when rem_static_win is True and value is set.""" + sample_params.rem_static_win = True + sample_params.static_threshold = 500 + + result = training_repo._get_static_threshold(sample_params) + + assert result == 500 + + def test_get_static_threshold_rem_static_win_true_with_none(self, training_repo, sample_params): + """Test returns 1 when rem_static_win is True and static_threshold is None.""" + sample_params.rem_static_win = True + sample_params.static_threshold = None + + result = training_repo._get_static_threshold(sample_params) + + assert result == 1 + + class TestInitScalerDict: """Tests for _init_scaler_dict method."""