From 0d85fcac2ce9bb2c156f57544d4c2b60134dbf0d Mon Sep 17 00:00:00 2001
From: Yoann Schneider <yschneider@teklia.com>
Date: Fri, 2 Sep 2022 10:15:38 +0000
Subject: [PATCH] Create missing types option when checking required types
existence
---
arkindex_worker/worker/element.py | 47 ++++++++++++++---
tests/test_elements_worker/test_elements.py | 57 ++++++++++++++++++---
2 files changed, 90 insertions(+), 14 deletions(-)
diff --git a/arkindex_worker/worker/element.py b/arkindex_worker/worker/element.py
index a22dbd0e..bf263e5c 100644
--- a/arkindex_worker/worker/element.py
+++ b/arkindex_worker/worker/element.py
@@ -4,13 +4,15 @@ ElementsWorker methods for elements and element types.
"""
import uuid
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Dict, Iterable, List, NamedTuple, Optional, Union
from peewee import IntegrityError
from arkindex_worker import logger
from arkindex_worker.cache import CachedElement, CachedImage
-from arkindex_worker.models import Element
+from arkindex_worker.models import Corpus, Element
+
+ElementType = NamedTuple("ElementType", name=str, slug=str, is_folder=bool)
class MissingTypeError(Exception):
@@ -24,13 +26,34 @@ class ElementMixin(object):
Mixin for the :class:`ElementsWorker` to provide ``Element`` helpers.
"""
- def check_required_types(self, corpus_id: str, *type_slugs: str) -> bool:
+ def create_required_types(self, corpus: Corpus, element_types: List[ElementType]):
+ """Creates given element types in the corpus.
+
+ :param Corpus corpus: The corpus to create types on.
+ :param List[ElementType] element_types: The missing element types to create.
+ """
+ for element_type in element_types:
+ self.request(
+ "CreateElementType",
+ body={
+ "slug": element_type.slug,
+ "display_name": element_type.name,
+ "folder": element_type.is_folder,
+ "corpus": corpus.id,
+ },
+ )
+ logger.info(f"Created a new element type with slug {element_type.slug}")
+
+ def check_required_types(
+ self, corpus_id: str, *type_slugs: str, create_missing: bool = False
+ ) -> bool:
"""
Check that a corpus has a list of required element types,
and raise an exception if any of them are missing.
:param str corpus_id: ID of the corpus to check types on.
:param str \\*type_slugs: Type slugs to look for.
+ :param bool create_missing: Whether missing types should be created.
:returns bool: True if all of the specified type slugs have been found.
:raises MissingTypeError: If any of the specified type slugs were not found.
"""
@@ -42,14 +65,22 @@ class ElementMixin(object):
isinstance(slug, str) for slug in type_slugs
), "Element type slugs must be strings."
- corpus = self.request("RetrieveCorpus", id=corpus_id)
- available_slugs = {element_type["slug"] for element_type in corpus["types"]}
+ corpus = Corpus(self.request("RetrieveCorpus", id=corpus_id))
+ available_slugs = {element_type.slug for element_type in corpus.types}
missing_slugs = set(type_slugs) - available_slugs
if missing_slugs:
- raise MissingTypeError(
- f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
- )
+ if create_missing:
+ self.create_required_types(
+ corpus,
+ element_types=[
+ ElementType(slug, slug, False) for slug in missing_slugs
+ ],
+ )
+ else:
+ raise MissingTypeError(
+ f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus.name} corpus ({corpus.id}).'
+ )
return True
diff --git a/tests/test_elements_worker/test_elements.py b/tests/test_elements_worker/test_elements.py
index 7ff16855..c74aa273 100644
--- a/tests/test_elements_worker/test_elements.py
+++ b/tests/test_elements_worker/test_elements.py
@@ -5,6 +5,7 @@ from uuid import UUID
import pytest
from apistar.exceptions import ErrorResponse
+from responses import matchers
from arkindex_worker.cache import (
SQL_VERSION,
@@ -33,11 +34,7 @@ def test_check_required_types_argument_types(mock_elements_worker):
assert str(e.value) == "Element type slugs must be strings."
-def test_check_required_types(monkeypatch, tmp_path, mock_elements_worker, responses):
- elements_path = tmp_path / "elements.json"
- elements_path.write_text("[]")
- monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
-
+def test_check_required_types(responses):
corpus_id = "12341234-1234-1234-1234-123412341234"
responses.add(
responses.GET,
@@ -49,7 +46,7 @@ def test_check_required_types(monkeypatch, tmp_path, mock_elements_worker, respo
},
)
worker = ElementsWorker()
- worker.configure()
+ worker.setup_api_client()
assert worker.check_required_types(corpus_id, "page")
assert worker.check_required_types(corpus_id, "page", "folder")
@@ -62,6 +59,54 @@ def test_check_required_types(monkeypatch, tmp_path, mock_elements_worker, respo
)
+def test_create_missing_types(responses):
+ corpus_id = "12341234-1234-1234-1234-123412341234"
+
+ responses.add(
+ responses.GET,
+ f"http://testserver/api/v1/corpus/{corpus_id}/",
+ json={
+ "id": corpus_id,
+ "name": "Some Corpus",
+ "types": [{"slug": "folder"}, {"slug": "page"}],
+ },
+ )
+ responses.add(
+ responses.POST,
+ "http://testserver/api/v1/elements/type/",
+ match=[
+ matchers.json_params_matcher(
+ {
+ "slug": "text_line",
+ "display_name": "text_line",
+ "folder": False,
+ "corpus": corpus_id,
+ }
+ )
+ ],
+ )
+ responses.add(
+ responses.POST,
+ "http://testserver/api/v1/elements/type/",
+ match=[
+ matchers.json_params_matcher(
+ {
+ "slug": "act",
+ "display_name": "act",
+ "folder": False,
+ "corpus": corpus_id,
+ }
+ )
+ ],
+ )
+ worker = ElementsWorker()
+ worker.setup_api_client()
+
+ assert worker.check_required_types(
+ corpus_id, "page", "text_line", "act", create_missing=True
+ )
+
+
def test_list_elements_elements_list_arg_wrong_type(
monkeypatch, tmp_path, mock_elements_worker
):
--
GitLab