From 33547faa9aedd5c49a8f2658ee9a3c647bc8a213 Mon Sep 17 00:00:00 2001 From: NolanB <nboukachab@teklia.com> Date: Fri, 2 Sep 2022 17:26:15 +0200 Subject: [PATCH] Replace ARKINDEX_CORPUS_ID .env variable with an attribute of base_worker corpus_id --- arkindex_worker/worker/base.py | 4 ++++ arkindex_worker/worker/classification.py | 5 +---- arkindex_worker/worker/element.py | 1 + arkindex_worker/worker/entity.py | 4 +--- tests/conftest.py | 2 -- tests/test_base_worker.py | 25 ++++++++++++++++----- tests/test_elements_worker/test_entities.py | 6 ++--- 7 files changed, 30 insertions(+), 17 deletions(-) diff --git a/arkindex_worker/worker/base.py b/arkindex_worker/worker/base.py index f0a571da..a893689c 100644 --- a/arkindex_worker/worker/base.py +++ b/arkindex_worker/worker/base.py @@ -119,6 +119,7 @@ class BaseWorker(object): logger.info(f"Worker will use {self.work_dir} as working directory") self.process_information = None + self.corpus_id = None self.user_configuration = {} self.support_cache = support_cache # use_cache will be updated in configure() if the cache is supported and if there @@ -188,6 +189,9 @@ class BaseWorker(object): # Load process information self.process_information = worker_run["process"] + # Load corpus id + self.corpus_id = worker_run["process"]["corpus"] + # Load worker version information worker_version = worker_run["worker_version"] self.worker_details = worker_version["worker"] diff --git a/arkindex_worker/worker/classification.py b/arkindex_worker/worker/classification.py index c9fdd086..5cdb63e6 100644 --- a/arkindex_worker/worker/classification.py +++ b/arkindex_worker/worker/classification.py @@ -2,9 +2,6 @@ """ ElementsWorker methods for classifications and ML classes. """ - -import os - from apistar.exceptions import ErrorResponse from peewee import IntegrityError @@ -45,7 +42,7 @@ class ClassificationMixin(object): :returns str: ID of the retrieved or created MLClass. """ if corpus_id is None: - corpus_id = os.environ.get("ARKINDEX_CORPUS_ID") + corpus_id = self.corpus_id if not self.classes.get(corpus_id): self.load_corpus_classes(corpus_id) diff --git a/arkindex_worker/worker/element.py b/arkindex_worker/worker/element.py index bf263e5c..daa53433 100644 --- a/arkindex_worker/worker/element.py +++ b/arkindex_worker/worker/element.py @@ -66,6 +66,7 @@ class ElementMixin(object): ), "Element type slugs must be strings." corpus = Corpus(self.request("RetrieveCorpus", id=corpus_id)) + # corpus = self.corpus_id available_slugs = {element_type.slug for element_type in corpus.types} missing_slugs = set(type_slugs) - available_slugs diff --git a/arkindex_worker/worker/entity.py b/arkindex_worker/worker/entity.py index 7b561c3e..b0b52903 100644 --- a/arkindex_worker/worker/entity.py +++ b/arkindex_worker/worker/entity.py @@ -2,8 +2,6 @@ """ ElementsWorker methods for entities. """ - -import os from enum import Enum from peewee import IntegrityError @@ -49,7 +47,7 @@ class EntityMixin(object): :type corpus: str or None """ if corpus is None: - corpus = os.environ.get("ARKINDEX_CORPUS_ID") + corpus = self.corpus_id assert element and isinstance( element, (Element, CachedElement) diff --git a/tests/conftest.py b/tests/conftest.py index f00ae5e7..70357545 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -106,7 +106,6 @@ def give_env_variable(request, monkeypatch): """Defines required environment variables""" monkeypatch.setenv("WORKER_VERSION_ID", "12341234-1234-1234-1234-123412341234") monkeypatch.setenv("ARKINDEX_WORKER_RUN_ID", "56785678-5678-5678-5678-567856785678") - monkeypatch.setenv("ARKINDEX_CORPUS_ID", "11111111-1111-1111-1111-111111111111") @pytest.fixture @@ -207,7 +206,6 @@ def mock_activity_calls(responses): def mock_elements_worker(monkeypatch, mock_worker_run_api): """Build and configure an ElementsWorker with fixed CLI parameters to avoid issues with pytest""" monkeypatch.setattr(sys, "argv", ["worker"]) - worker = ElementsWorker() worker.configure() return worker diff --git a/tests/test_base_worker.py b/tests/test_base_worker.py index a3262ee9..90b89736 100644 --- a/tests/test_base_worker.py +++ b/tests/test_base_worker.py @@ -194,7 +194,10 @@ def test_configure_worker_run(mocker, monkeypatch, responses): "configuration": {"configuration": {}}, }, "configuration": user_configuration, - "process": {"id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"}, + "process": { + "id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff", + "corpus": "11111111-1111-1111-1111-111111111111", + }, } responses.add( @@ -257,7 +260,10 @@ def test_configure_user_configuration_defaults( "param_5": True, }, }, - "process": {"id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"}, + "process": { + "id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff", + "corpus": "11111111-1111-1111-1111-111111111111", + }, } responses.add( responses.GET, @@ -310,7 +316,10 @@ def test_configure_user_config_debug(mocker, monkeypatch, responses, debug): "name": "BBB", "configuration": {"debug": debug}, }, - "process": {"id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"}, + "process": { + "id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff", + "corpus": "11111111-1111-1111-1111-111111111111", + }, } responses.add( responses.GET, @@ -357,7 +366,10 @@ def test_configure_worker_run_missing_conf(mocker, monkeypatch, responses): "configuration": {"configuration": {}}, }, "configuration": {"id": "bbbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb", "name": "BBB"}, - "process": {"id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"}, + "process": { + "id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff", + "corpus": "11111111-1111-1111-1111-111111111111", + }, } responses.add( responses.GET, @@ -404,7 +416,10 @@ def test_configure_worker_run_no_worker_run_conf(mocker, monkeypatch, responses) "configuration": {}, }, "configuration": None, - "process": {"id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"}, + "process": { + "id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff", + "corpus": "11111111-1111-1111-1111-111111111111", + }, } responses.add( responses.GET, diff --git a/tests/test_elements_worker/test_entities.py b/tests/test_elements_worker/test_entities.py index 57ea05af..71ed34ab 100644 --- a/tests/test_elements_worker/test_entities.py +++ b/tests/test_elements_worker/test_entities.py @@ -101,7 +101,6 @@ 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, @@ -111,8 +110,8 @@ def test_create_entity_wrong_corpus(monkeypatch, mock_elements_worker): ) 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") + # # Removing corpus_id variable should give an error when corpus=None + mock_elements_worker.corpus_id = None with pytest.raises(AssertionError) as e: mock_elements_worker.create_entity( element=elt, @@ -130,6 +129,7 @@ def test_create_entity_wrong_corpus(monkeypatch, mock_elements_worker): corpus=1234, ) assert str(e.value) == "corpus shouldn't be null and should be of type str" + mock_elements_worker.corpus_id = ("11111111-1111-1111-1111-111111111111",) def test_create_entity_wrong_metas(mock_elements_worker): -- GitLab