Use ARKINDEX_CORPUS_ID env variable in get_ml_class_id

arkindex_worker/worker/classification.py

 # -*- coding: utf-8 -*-
+import os
+
 from apistar.exceptions import ErrorResponse
 
 from arkindex_worker import logger
@@ -24,6 +26,9 @@ class ClassificationMixin(object):
         Return the ID corresponding to the given class name on a specific corpus
         This method will automatically create missing classes
         """
+        if corpus_id is None:
+            corpus_id = os.environ.get("ARKINDEX_CORPUS_ID")
+
         if not self.classes.get(corpus_id):
             self.load_corpus_classes(corpus_id)
 
@@ -82,7 +87,7 @@ class ClassificationMixin(object):
                 "CreateClassification",
                 body={
                     "element": element.id,
-                    "ml_class": self.get_ml_class_id(element.corpus.id, ml_class),
+                    "ml_class": self.get_ml_class_id(None, ml_class),
                     "worker_version": self.worker_version_id,
                     "confidence": confidence,
                     "high_confidence": high_confidence,

tests/conftest.py

@@ -165,6 +165,7 @@ def mock_user_api(responses):
 def mock_elements_worker(monkeypatch, mock_worker_version_api):
     """Build and configure an ElementsWorker with fixed CLI parameters to avoid issues with pytest"""
     monkeypatch.setattr(sys, "argv", ["worker"])
+    monkeypatch.setenv("ARKINDEX_CORPUS_ID", "11111111-1111-1111-1111-111111111111")
 
     worker = ElementsWorker()
     worker.configure()
@@ -185,6 +186,7 @@ def mock_base_worker_with_cache(mocker, monkeypatch, mock_worker_version_api):
 def mock_elements_worker_with_cache(monkeypatch, mock_worker_version_api):
     """Build and configure an ElementsWorker using SQLite cache with fixed CLI parameters to avoid issues with pytest"""
     monkeypatch.setattr(sys, "argv", ["worker"])
+    monkeypatch.setenv("ARKINDEX_CORPUS_ID", "11111111-1111-1111-1111-111111111111")
 
     worker = ElementsWorker(use_cache=True)
     worker.configure()

tests/test_elements_worker/test_classifications.py

@@ -172,12 +172,7 @@ def test_create_classification_wrong_element(mock_elements_worker):
 
 
 def test_create_classification_wrong_ml_class(mock_elements_worker, responses):
-    elt = Element(
-        {
-            "id": "12341234-1234-1234-1234-123412341234",
-            "corpus": {"id": "11111111-1111-1111-1111-111111111111"},
-        }
-    )
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
 
     with pytest.raises(AssertionError) as e:
         mock_elements_worker.create_classification(
@@ -249,12 +244,7 @@ def test_create_classification_wrong_confidence(mock_elements_worker):
     mock_elements_worker.classes = {
         "11111111-1111-1111-1111-111111111111": {"a_class": "0000"}
     }
-    elt = Element(
-        {
-            "id": "12341234-1234-1234-1234-123412341234",
-            "corpus": {"id": "11111111-1111-1111-1111-111111111111"},
-        }
-    )
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
     with pytest.raises(AssertionError) as e:
         mock_elements_worker.create_classification(
             element=elt,
@@ -308,12 +298,7 @@ def test_create_classification_wrong_high_confidence(mock_elements_worker):
     mock_elements_worker.classes = {
         "11111111-1111-1111-1111-111111111111": {"a_class": "0000"}
     }
-    elt = Element(
-        {
-            "id": "12341234-1234-1234-1234-123412341234",
-            "corpus": {"id": "11111111-1111-1111-1111-111111111111"},
-        }
-    )
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
 
     with pytest.raises(AssertionError) as e:
         mock_elements_worker.create_classification(
@@ -342,12 +327,7 @@ def test_create_classification_api_error(responses, mock_elements_worker):
     mock_elements_worker.classes = {
         "11111111-1111-1111-1111-111111111111": {"a_class": "0000"}
     }
-    elt = Element(
-        {
-            "id": "12341234-1234-1234-1234-123412341234",
-            "corpus": {"id": "11111111-1111-1111-1111-111111111111"},
-        }
-    )
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
     responses.add(
         responses.POST,
         "http://testserver/api/v1/classifications/",
@@ -379,12 +359,7 @@ def test_create_classification(responses, mock_elements_worker):
     mock_elements_worker.classes = {
         "11111111-1111-1111-1111-111111111111": {"a_class": "0000"}
     }
-    elt = Element(
-        {
-            "id": "12341234-1234-1234-1234-123412341234",
-            "corpus": {"id": "11111111-1111-1111-1111-111111111111"},
-        }
-    )
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
     responses.add(
         responses.POST,
         "http://testserver/api/v1/classifications/",
@@ -423,12 +398,7 @@ def test_create_classification_duplicate(responses, mock_elements_worker):
     mock_elements_worker.classes = {
         "11111111-1111-1111-1111-111111111111": {"a_class": "0000"}
     }
-    elt = Element(
-        {
-            "id": "12341234-1234-1234-1234-123412341234",
-            "corpus": {"id": "11111111-1111-1111-1111-111111111111"},
-        }
-    )
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
     responses.add(
         responses.POST,
         "http://testserver/api/v1/classifications/",