Support giving a CachedElement to create_entity function

arkindex_worker/worker/entity.py

@@ -4,6 +4,7 @@
 from enum import Enum
 
 from arkindex_worker import logger
+from arkindex_worker.cache import CachedElement
 from arkindex_worker.models import Element
 
 
@@ -24,8 +25,8 @@ class EntityMixin(object):
         Return the ID of the created entity
         """
         assert element and isinstance(
-            element, Element
-        ), "element shouldn't be null and should be of type Element"
+            element, (Element, CachedElement)
+        ), "element shouldn't be null and should be an Element or CachedElement"
         assert name and isinstance(
             name, str
         ), "name shouldn't be null and should be of type str"

tests/test_elements_worker/test_entities.py

@@ -16,7 +16,10 @@ def test_create_entity_wrong_element(mock_elements_worker):
             type=EntityType.Person,
             corpus="12341234-1234-1234-1234-123412341234",
         )
-    assert str(e.value) == "element shouldn't be null and should be of type Element"
+    assert (
+        str(e.value)
+        == "element shouldn't be null and should be an Element or CachedElement"
+    )
 
     with pytest.raises(AssertionError) as e:
         mock_elements_worker.create_entity(
@@ -25,7 +28,10 @@ def test_create_entity_wrong_element(mock_elements_worker):
             type=EntityType.Person,
             corpus="12341234-1234-1234-1234-123412341234",
         )
-    assert str(e.value) == "element shouldn't be null and should be of type Element"
+    assert (
+        str(e.value)
+        == "element shouldn't be null and should be an Element or CachedElement"
+    )
 
 
 def test_create_entity_wrong_name(mock_elements_worker):