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.
This commit is contained in:
@@ -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')
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user