diff --git a/arkindex_worker/worker/entity.py b/arkindex_worker/worker/entity.py index b29b267c7ab5b6193f9fb9a46be831bb6109e32c..30c1637d82037ec0cd72e0cf27e0de3d29000e19 100644 --- a/arkindex_worker/worker/entity.py +++ b/arkindex_worker/worker/entity.py @@ -1,6 +1,5 @@ # -*- coding: utf-8 -*- - - +import os from enum import Enum from arkindex_worker import logger @@ -19,11 +18,16 @@ class EntityType(Enum): class EntityMixin(object): - def create_entity(self, element, name, type, corpus, metas=None, validated=None): + def create_entity( + self, element, name, type, corpus=None, metas=None, validated=None + ): """ Create an entity on the given corpus through API Return the ID of the created entity """ + if corpus is None: + corpus = os.environ.get("ARKINDEX_CORPUS_ID") + assert element and isinstance( element, (Element, CachedElement) ), "element shouldn't be null and should be an Element or CachedElement" diff --git a/tests/test_elements_worker/test_entities.py b/tests/test_elements_worker/test_entities.py index caa0a34193b9a63e674b6ac228acf14a48397fd8..1fe38b36ba4fc983f09f3a4c813112d666767df7 100644 --- a/tests/test_elements_worker/test_entities.py +++ b/tests/test_elements_worker/test_entities.py @@ -87,9 +87,22 @@ def test_create_entity_wrong_type(mock_elements_worker): assert str(e.value) == "type shouldn't be null and should be of type EntityType" -def test_create_entity_wrong_corpus(mock_elements_worker): +def test_create_entity_wrong_corpus(monkeypatch, mock_elements_worker): elt = Element({"id": "12341234-1234-1234-1234-123412341234"}) + # Triggering an error on metas param, not giving corpus should work since + # ARKINDEX_CORPUS_ID environment variable is set on mock_elements_worker + with pytest.raises(AssertionError) as e: + mock_elements_worker.create_entity( + element=elt, + name="Bob Bob", + type=EntityType.Person, + metas="wrong metas", + ) + assert str(e.value) == "metas should be of type dict" + + # Removing ARKINDEX_CORPUS_ID environment variable should give an error when corpus=None + monkeypatch.delenv("ARKINDEX_CORPUS_ID") with pytest.raises(AssertionError) as e: mock_elements_worker.create_entity( element=elt,