diff --git a/tests/test_ilastik.py b/tests/test_ilastik.py
index 8b687874f85e1b0ae65a17dcc569da5e3e0a3379..01ac4306622f8668a2581d04d3c472f29aa332d8 100644
--- a/tests/test_ilastik.py
+++ b/tests/test_ilastik.py
@@ -110,6 +110,9 @@ class TestIlastikOverApi(TestServerBaseClass):
         rj = resp_list.json()
         self.assertEqual(rj[model_id]['class'], 'IlastikPixelClassifierModel')
 
+        return model_id
+
+
     def test_load_ilastik_object_model(self):
         resp_load = requests.put(
             self.uri + 'models/ilastik/object_classification/load/',
@@ -121,4 +124,18 @@ class TestIlastikOverApi(TestServerBaseClass):
         resp_list = requests.get(self.uri + 'models')
         self.assertEqual(resp_list.status_code, 200)
         rj = resp_list.json()
-        self.assertEqual(rj[model_id]['class'], 'IlastikObjectClassifierModel')
\ No newline at end of file
+        self.assertEqual(rj[model_id]['class'], 'IlastikObjectClassifierModel')
+
+    def test_ilastik_infer_pixel_probability(self):
+        TestServerBaseClass.copy_input_file_to_server()
+        model_id = self.test_load_ilastik_pixel_model()
+
+        resp_infer = requests.put(
+            self.uri + f'infer/from_image_file',
+            params={
+                'model_id': model_id,
+                'input_filename': czifile['filename'],
+                'channel': 2,
+            },
+        )
+        self.assertEqual(resp_infer.status_code, 200, resp_infer.content.decode())
diff --git a/tests/test_model.py b/tests/test_model.py
index ca0324a9e574447cf52d591ac2edc0bafa03829c..a265cbe4bd9860b8b053ca9ad768a7ca41885554 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1,7 +1,7 @@
 import unittest
 from conf.testing import czifile
 from model_server.image import CziImageFileAccessor
-from model_server.model import DummyImageToImageModel, CouldNotLoadModelError, Model
+from model_server.model import DummyImageToImageModel, CouldNotLoadModelError
 
 class TestCziImageFileAccess(unittest.TestCase):
     def setUp(self) -> None: