diff --git a/arkindex_worker/worker/version.py b/arkindex_worker/worker/version.py index f32fdc3a2f3dc1780113bcf3d451e15841a88aee..39d75aabe2bd22624cdcf161fa82c9155b309e50 100644 --- a/arkindex_worker/worker/version.py +++ b/arkindex_worker/worker/version.py @@ -14,7 +14,6 @@ class WorkerVersionMixin(object): return worker_version - def get_ml_result_slug(self, ml_result) -> str: """ Helper function to get the worker slug from element, classification or transcription. @@ -24,13 +23,13 @@ class WorkerVersionMixin(object): """ # Handle cached models - if hasattr(ml_result, 'worker_version_id'): + if hasattr(ml_result, "worker_version_id"): worker_version_id = ml_result.worker_version_id elif "worker_version" in ml_result: - worker_version_id = ml_result['worker_version'] + worker_version_id = ml_result["worker_version"] # transcriptions have worker_version_id but elements have worker_version elif "worker_version_id" in ml_result: - worker_version_id = ml_result['worker_version_id'] + worker_version_id = ml_result["worker_version_id"] else: raise ValueError(f"Unable to get slug from: {ml_result}") @@ -38,4 +37,4 @@ class WorkerVersionMixin(object): return worker_version = self.get_worker_version(worker_version_id) - return worker_version['worker']['slug'] + return worker_version["worker"]["slug"] diff --git a/tests/test_elements_worker/test_worker.py b/tests/test_elements_worker/test_worker.py index 6c520a9dc8704ab614ec0d280d3a8495abe8d069..ce85d12dd14b8cff0ec4a40d434c87c221a79e3d 100644 --- a/tests/test_elements_worker/test_worker.py +++ b/tests/test_elements_worker/test_worker.py @@ -4,8 +4,8 @@ import json import pytest from apistar.exceptions import ErrorResponse -from arkindex_worker.models import Element from arkindex_worker.cache import CachedElement +from arkindex_worker.models import Element from arkindex_worker.worker import ActivityState # Common API calls for all workers @@ -123,9 +123,7 @@ def test_get_ml_result_slug__ok(mocker, fake_dummy_worker, ml_result, expected_s fake_dummy_worker.get_worker_version = mocker.MagicMock() fake_dummy_worker.get_worker_version.return_value = { "id": TEST_VERSION_ID, - "worker": { - "slug": "mock_slug" - } + "worker": {"slug": "mock_slug"}, } slug = fake_dummy_worker.get_ml_result_slug(ml_result)