SIENTIAPDE-1430: Refactor static threshold calculation logic and enhance test coverage
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user