Cache entities in the worker when calling the `list_corpus_entities` helper

arkindex_worker/worker/entity.py

@@ -331,8 +331,7 @@ class EntityMixin:
         parent: Element | None = None,
     ):
         """
-        List all entities in the worker's corpus
-        This method does not support cache
+        List all entities in the worker's corpus and store them in the ``self.entities`` cache.
         :param name: Filter entities by part of their name (case-insensitive)
         :param parent: Restrict entities to those linked to all transcriptions of an element and all its descendants. Note that links to metadata are ignored.
         """
@@ -346,8 +345,14 @@ class EntityMixin:
             assert isinstance(parent, Element), "parent should be of type Element"
             query_params["parent"] = parent.id
 
-        return self.api_client.paginate(
-            "ListCorpusEntities", id=self.corpus_id, **query_params
+        self.entities = {
+            entity["id"]: entity
+            for entity in self.api_client.paginate(
+                "ListCorpusEntities", id=self.corpus_id, **query_params
+            )
+        }
+        logger.info(
+            f"Loaded {len(self.entities)} entities in corpus ({self.corpus_id})"
         )
 
     def list_corpus_entity_types(

tests/test_elements_worker/test_entities.py

@@ -741,12 +741,13 @@ def test_list_corpus_entities(responses, mock_elements_worker):
         },
     )
 
-    # list is required to actually do the request
-    assert list(mock_elements_worker.list_corpus_entities()) == [
-        {
+    mock_elements_worker.list_corpus_entities()
+
+    assert mock_elements_worker.entities == {
+        "fake_entity_id": {
             "id": "fake_entity_id",
         }
-    ]
+    }
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [