Fix tests

deepsensor/model/model.py

@@ -510,8 +510,8 @@ def unnormalise_pred_array(arr, **kwargs):
                 try:
                     method = getattr(self, param)
                     prediction_methods[param] = method
-                except ValueError:
-                    raise ValueError(
+                except AttributeError:
+                    raise AttributeError(
                         f"Prediction method {param} not found in model class."
                     )
             if n_samples >= 1:

tests/test_model.py

@@ -448,10 +448,10 @@ def test_highlevel_predict_with_pred_params(self):
         pred_params = ["mean", "std", "variance"]
         pred = model.predict(task, X_t=self.da, pred_params=pred_params)
         for pred_param in pred_params:
-            assert pred_param in pred["air"]
+            assert pred_param in pred["var"]
 
-        # Check that passing an invalid parameter raises a ValueError
-        with self.assertRaises(ValueError):
+        # Check that passing an invalid parameter raises an AttributeError
+        with self.assertRaises(AttributeError):
             model.predict(task, X_t=self.da, pred_params=["invalid_param"])
 
     def test_saving_and_loading(self):