diff --git a/arkindex_worker/worker/classification.py b/arkindex_worker/worker/classification.py index d9dfaad1dadead0e1843172bb373e17ab2c8653c..3e979c043cceac8d847fd99d291797ee8e74e5cb 100644 --- a/arkindex_worker/worker/classification.py +++ b/arkindex_worker/worker/classification.py @@ -1,4 +1,6 @@ # -*- 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, diff --git a/tests/conftest.py b/tests/conftest.py index 7141bad3fd94235721438f4580b31e5b2280baa6..265f5d05c7b5cfbc816b07a0b14f72eb994b157d 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_elements_worker/test_classifications.py b/tests/test_elements_worker/test_classifications.py index 05e8d8d615a6cffb8de9ca3c3bd99548a14dcca1..20a677071451b76cc329816fd91787f7bcd571c0 100644 --- a/tests/test_elements_worker/test_classifications.py +++ b/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/",