From 6890ca2af354f5002b19023d54afbc80c095c524 Mon Sep 17 00:00:00 2001
From: Eva Bardou <ebardou@teklia.com>
Date: Thu, 22 Apr 2021 10:40:53 +0200
Subject: [PATCH] Retrieve ARKINDEX_CORPUS_ID in create_entity when corpus is
 None

---
 arkindex_worker/worker/entity.py            |  6 ++++--
 tests/test_elements_worker/test_entities.py | 16 +++++++++++++++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/arkindex_worker/worker/entity.py b/arkindex_worker/worker/entity.py
index b29b267c..e42036e1 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
@@ -24,6 +23,9 @@ class EntityMixin(object):
         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 caa0a341..540635de 100644
--- a/tests/test_elements_worker/test_entities.py
+++ b/tests/test_elements_worker/test_entities.py
@@ -87,9 +87,23 @@ 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, giving corpus=None 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,
+            corpus=None,
+            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,
-- 
GitLab