diff --git a/test/unit/repositories/test_image.py b/test/unit/repositories/test_image.py
index de55b6b..6dacb41 100644
--- a/test/unit/repositories/test_image.py
+++ b/test/unit/repositories/test_image.py
@@ -116,9 +116,9 @@ def call_list(self, entity_id: Optional[str], primary: Optional[bool]) -> None:
             session=self.mock_session, entity_id=entity_id, primary=primary
         )
 
-    def check_list_success(self) -> None:
+    def check_list_success(self, query: Optional[dict] = None) -> None:
         """Checks that a prior call to `call_list` worked as expected."""
-        self.images_collection.find.assert_called_once_with(session=self.mock_session)
+        self.images_collection.find.assert_called_once_with(session=self.mock_session, query=query)
         assert self._obtained_image_out == self._expected_image_out
 
 
@@ -142,18 +142,18 @@ def test_list(self):
         query = {"entity_id": IMAGE_IN_DATA_ALL_VALUES["entity_id"]}
         self.mock_list([IMAGE_IN_DATA_ALL_VALUES], query=query)
         self.call_list(entity_id=query["entity_id"])
-        self.check_list_success()
+        self.check_list_success(query)
 
     def test_list(self):
         """Test listing all images with a primary argument."""
         query = {"primary": False}
         self.mock_list([IMAGE_IN_DATA_ALL_VALUES], query=query)
         self.call_list(primary=query["primary"])
-        self.check_list_success()
+        self.check_list_success(query)
 
     def test_list(self):
         """Test listing all images with an entity_id and primary argument."""
         query = {"entity_id": IMAGE_IN_DATA_ALL_VALUES["entity_id"], "primary": False}
         self.mock_list([IMAGE_IN_DATA_ALL_VALUES], query=query)
         self.call_list(primary=query["primary"], entity_id=query["entity_id"])
-        self.check_list_success()
+        self.check_list_success(query)