From c5fb921654700635bf924aa175b43829f0d7e18c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 21 Oct 2025 15:58:28 -0300 Subject: [PATCH] SIENTIAPDE-1312 Refactor error handling in MLFlowRepository and update tests - Improved error message formatting in MLFlowRepository for better readability. - Updated test assertions to ensure correct calls to MLflow methods during experiment retrieval and creation. --- laborious/utils/repository/model_repository.py | 4 +++- tests/laborious/utils/repository/test_model_repository.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 1ab72c1..7400502 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -178,7 +178,9 @@ class MLFlowRepository: experiment_id = mlflow.create_experiment(experiment_name) experiment = mlflow.get_experiment(experiment_id) if experiment is None: - raise ValueError(f'Experiment {experiment_name} not found after creation, unknown reason') + raise ValueError( + f'Experiment {experiment_name} not found after creation, unknown reason' + ) else: raise ValueError(f'Experiment {experiment_name} not found') diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index a437653..c21608d 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -157,10 +157,13 @@ def test_get_experiment_none_create(mlflow, mlflow_repository): mlflow.get_experiment_by_name.return_value = None - mlflow.create_experiment.return_value = experiment + mlflow.get_experiment.return_value = experiment output = mlflow_repository.get_experiment('test', create_if_not_exists=True) + mlflow.create_experiment.assert_called_once_with('test') + mlflow.get_experiment.assert_called_once_with(mlflow.create_experiment.return_value) + assert output == experiment