diff --git a/src/unittests/java/MakeItFit/MakeItFitTest.java b/src/unittests/java/MakeItFit/MakeItFitTest.java
index 9af3bd3..f375a1a 100644
--- a/src/unittests/java/MakeItFit/MakeItFitTest.java
+++ b/src/unittests/java/MakeItFit/MakeItFitTest.java
@@ -413,7 +413,7 @@ void testRemoveTrainingPlan() throws Exception {
         UUID          planCode  = model.createTrainingPlan(userCode, startDate);
 
         model.removeTrainingPlan(planCode);
-        assertThrows(IllegalArgumentException.class, () -> { model.getTrainingPlan(planCode); });
+        assertThrows(EntityDoesNotExistException.class, () -> { model.getTrainingPlan(planCode); });
     }
 
     @Test
@@ -443,8 +443,7 @@ void testGetTrainingPlan() throws Exception {
     @Test
     void testGetTrainingPlanException() {
         UUID invalidCode = UUID.randomUUID();
-
-        assertThrows(IllegalArgumentException.class, () -> { model.getTrainingPlan(invalidCode); });
+        assertThrows(EntityDoesNotExistException.class, () -> { model.getTrainingPlan(invalidCode); });
     }
 
     @Test