diff --git a/arkindex_worker/worker/entity.py b/arkindex_worker/worker/entity.py index 81a314571d2742a823f375736a82f7f59fe3d1bd..b29b267c7ab5b6193f9fb9a46be831bb6109e32c 100644 --- a/arkindex_worker/worker/entity.py +++ b/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" diff --git a/tests/test_elements_worker/test_entities.py b/tests/test_elements_worker/test_entities.py index be513d6b6a7ad2e1f21baf2099fdb3f7dd4de855..caa0a34193b9a63e674b6ac228acf14a48397fd8 100644 --- a/tests/test_elements_worker/test_entities.py +++ b/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):