diff --git a/test/ui/test_ui.py b/test/ui/test_ui.py
index 7efb442..f153cff 100644
--- a/test/ui/test_ui.py
+++ b/test/ui/test_ui.py
@@ -43,6 +43,7 @@ def test_user_calibrated_models(petab_select_problem):
     assert model_M1_0.model_subspace_id == "M1_0"
     model_M1_0.set_criterion(petab_select_problem.criterion, 100)
     iteration_results = petab_select.ui.end_iteration(
+        problem=petab_select_problem,
         candidate_space=iteration[CANDIDATE_SPACE],
         calibrated_models=[model_M1_0],
     )
@@ -61,6 +62,7 @@ def test_user_calibrated_models(petab_select_problem):
     for uncalibrated_model in iteration[UNCALIBRATED_MODELS]:
         uncalibrated_model.set_criterion(petab_select_problem.criterion, 50)
     iteration_results = petab_select.ui.end_iteration(
+        problem=petab_select_problem,
         candidate_space=iteration[CANDIDATE_SPACE],
         calibrated_models=iteration[UNCALIBRATED_MODELS],
     )