arkindex-base-worker 0.4.0a1__tar.gz → 0.4.0b1__tar.gz

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.4.0a1
+Version: 0.4.0b1
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>
@@ -41,17 +41,17 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: peewee~=3.17
-Requires-Dist: Pillow==10.3.0
+Requires-Dist: Pillow==10.4.0
 Requires-Dist: python-gnupg==0.5.2
-Requires-Dist: shapely==2.0.3
+Requires-Dist: shapely==2.0.5
 Requires-Dist: teklia-toolbox==0.1.5
 Requires-Dist: zstandard==0.22.0
 Provides-Extra: docs
-Requires-Dist: black==24.4.0; extra == "docs"
-Requires-Dist: mkdocs-material==9.5.17; extra == "docs"
-Requires-Dist: mkdocstrings-python==1.9.2; extra == "docs"
+Requires-Dist: black==24.4.2; extra == "docs"
+Requires-Dist: mkdocs-material==9.5.31; extra == "docs"
+Requires-Dist: mkdocstrings-python==1.10.7; extra == "docs"
 Provides-Extra: tests
-Requires-Dist: pytest==8.1.1; extra == "tests"
+Requires-Dist: pytest==8.3.2; extra == "tests"
 Requires-Dist: pytest-mock==3.14.0; extra == "tests"
 Requires-Dist: pytest-responses==0.5.1; extra == "tests"
 

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_base_worker.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.4.0a1
+Version: 0.4.0b1
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>
@@ -41,17 +41,17 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: peewee~=3.17
-Requires-Dist: Pillow==10.3.0
+Requires-Dist: Pillow==10.4.0
 Requires-Dist: python-gnupg==0.5.2
-Requires-Dist: shapely==2.0.3
+Requires-Dist: shapely==2.0.5
 Requires-Dist: teklia-toolbox==0.1.5
 Requires-Dist: zstandard==0.22.0
 Provides-Extra: docs
-Requires-Dist: black==24.4.0; extra == "docs"
-Requires-Dist: mkdocs-material==9.5.17; extra == "docs"
-Requires-Dist: mkdocstrings-python==1.9.2; extra == "docs"
+Requires-Dist: black==24.4.2; extra == "docs"
+Requires-Dist: mkdocs-material==9.5.31; extra == "docs"
+Requires-Dist: mkdocstrings-python==1.10.7; extra == "docs"
 Provides-Extra: tests
-Requires-Dist: pytest==8.1.1; extra == "tests"
+Requires-Dist: pytest==8.3.2; extra == "tests"
 Requires-Dist: pytest-mock==3.14.0; extra == "tests"
 Requires-Dist: pytest-responses==0.5.1; extra == "tests"
 

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_base_worker.egg-info/requires.txt

@@ -1,16 +1,16 @@
 peewee~=3.17
-Pillow==10.3.0
+Pillow==10.4.0
 python-gnupg==0.5.2
-shapely==2.0.3
+shapely==2.0.5
 teklia-toolbox==0.1.5
 zstandard==0.22.0
 
 [docs]
-black==24.4.0
-mkdocs-material==9.5.17
-mkdocstrings-python==1.9.2
+black==24.4.2
+mkdocs-material==9.5.31
+mkdocstrings-python==1.10.7
 
 [tests]
-pytest==8.1.1
+pytest==8.3.2
 pytest-mock==3.14.0
 pytest-responses==0.5.1

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/__init__.py

@@ -83,7 +83,20 @@ class ElementsWorker(
         """
         super().__init__(description, support_cache)
 
-        # Add mandatory argument to process elements
+        self.classes = {}
+
+        self.entity_types = {}
+        """Known and available entity types in processed corpus
+        """
+
+        self.corpus_types = {}
+        """Known and available element types in processed corpus
+        """
+
+        self._worker_version_cache = {}
+
+    def add_arguments(self):
+        """Define specific ``argparse`` arguments for this worker"""
         self.parser.add_argument(
             "--elements-list",
             help="JSON elements list to use",
@@ -97,14 +110,6 @@ class ElementsWorker(
             help="One or more Arkindex element ID",
         )
 
-        self.classes = {}
-
-        self.entity_types = {}
-        """Known and available entity types in processed corpus
-        """
-
-        self._worker_version_cache = {}
-
     def list_elements(self) -> Iterable[CachedElement] | list[str]:
         """
         List the elements to be processed, either from the CLI arguments or
@@ -222,7 +227,9 @@ class ElementsWorker(
                     element = item
                 else:
                     # Load element using the Arkindex API
-                    element = Element(**self.request("RetrieveElement", id=item))
+                    element = Element(
+                        **self.api_client.request("RetrieveElement", id=item)
+                    )
 
                 logger.info(f"Processing {element} ({i}/{count})")
 
@@ -301,7 +308,7 @@ class ElementsWorker(
         assert isinstance(state, ActivityState), "state should be an ActivityState"
 
         try:
-            self.request(
+            self.api_client.request(
                 "UpdateWorkerActivity",
                 id=self.worker_run_id,
                 body={
@@ -376,6 +383,8 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         # Set as an instance variable as dataset workers might use it to easily extract its content
         self.downloaded_dataset_artifact: Path | None = None
 
+    def add_arguments(self):
+        """Define specific ``argparse`` arguments for this worker"""
         self.parser.add_argument(
             "--set",
             type=check_dataset_set,
@@ -472,7 +481,7 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
             # Retrieving dataset information is not already cached
             if dataset_id not in datasets:
                 datasets[dataset_id] = Dataset(
-                    **self.request("RetrieveDataset", id=dataset_id)
+                    **self.api_client.request("RetrieveDataset", id=dataset_id)
                 )
 
             yield Set(name=set_name, dataset=datasets[dataset_id])

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/base.py

@@ -231,7 +231,7 @@ class BaseWorker:
             logger.debug("Debug output enabled")
 
         # Load worker run information
-        worker_run = self.request("RetrieveWorkerRun", id=self.worker_run_id)
+        worker_run = self.api_client.request("RetrieveWorkerRun", id=self.worker_run_id)
 
         # Load process information
         self.process_information = worker_run["process"]
@@ -290,7 +290,7 @@ class BaseWorker:
         if self.support_cache and self.args.database is not None:
             self.use_cache = True
         elif self.support_cache and self.task_id:
-            task = self.request("RetrieveTaskFromAgent", id=self.task_id)
+            task = self.api_client.request("RetrieveTask", id=self.task_id)
             self.task_parents = task["parents"]
             paths = self.find_parents_file_paths(Path("db.sqlite"))
             self.use_cache = len(paths) > 0
@@ -331,7 +331,7 @@ class BaseWorker:
 
         # Load from the backend
         try:
-            resp = self.request("RetrieveSecret", name=str(name))
+            resp = self.api_client.request("RetrieveSecret", name=str(name))
             secret = resp["content"]
             logging.info(f"Loaded API secret {name}")
         except ErrorResponse as e:
@@ -471,12 +471,6 @@ class BaseWorker:
             # Clean up
             shutil.rmtree(base_extracted_path)
 
-    def request(self, *args, **kwargs):
-        """
-        Wrapper around the ``ArkindexClient.request`` method.
-        """
-        return self.api_client.request(*args, **kwargs)
-
     def add_arguments(self):
         """Override this method to add ``argparse`` arguments to this worker"""
 

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/classification.py

@@ -39,7 +39,7 @@ class ClassificationMixin:
         if ml_class_id is None:
             logger.info(f"Creating ML class {ml_class} on corpus {self.corpus_id}")
             try:
-                response = self.request(
+                response = self.api_client.request(
                     "CreateMLClass", id=self.corpus_id, body={"name": ml_class}
                 )
                 ml_class_id = self.classes[ml_class] = response["id"]
@@ -119,7 +119,7 @@ class ClassificationMixin:
             )
             return
         try:
-            created = self.request(
+            created = self.api_client.request(
                 "CreateClassification",
                 body={
                     "element": str(element.id),
@@ -220,7 +220,7 @@ class ClassificationMixin:
             )
             return
 
-        created_cls = self.request(
+        created_cls = self.api_client.request(
             "CreateClassifications",
             body={
                 "parent": str(element.id),

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/corpus.py

@@ -63,7 +63,9 @@ class CorpusMixin:
         # Download latest export
         export_id: str = exports[0]["id"]
         logger.info(f"Downloading export ({export_id})...")
-        export: _TemporaryFileWrapper = self.request("DownloadExport", id=export_id)
+        export: _TemporaryFileWrapper = self.api_client.request(
+            "DownloadExport", id=export_id
+        )
         logger.info(f"Downloaded export ({export_id}) @ `{export.name}`")
 
         return export

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/dataset.py

@@ -93,7 +93,7 @@ class DatasetMixin:
             logger.warning("Cannot update dataset as this worker is in read-only mode")
             return
 
-        updated_dataset = self.request(
+        updated_dataset = self.api_client.request(
             "PartialUpdateDataset",
             id=dataset.id,
             body={"state": state.value},

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/element.py

@@ -31,6 +31,21 @@ class MissingTypeError(Exception):
 
 
 class ElementMixin:
+    def list_corpus_types(self):
+        """
+        Loads available element types in corpus.
+        """
+        self.corpus_types = {
+            element_type["slug"]: element_type
+            for element_type in self.api_client.request(
+                "RetrieveCorpus", id=self.corpus_id
+            )["types"]
+        }
+        count = len(self.corpus_types)
+        logger.info(
+            f'Loaded {count} element type{"s"[:count>1]} in corpus ({self.corpus_id}).'
+        )
+
     @unsupported_cache
     def create_required_types(self, element_types: list[ElementType]):
         """Creates given element types in the corpus.
@@ -38,7 +53,7 @@ class ElementMixin:
         :param element_types: The missing element types to create.
         """
         for element_type in element_types:
-            self.request(
+            self.api_client.request(
                 "CreateElementType",
                 body={
                     "slug": element_type.slug,
@@ -66,10 +81,10 @@ class ElementMixin:
             isinstance(slug, str) for slug in type_slugs
         ), "Element type slugs must be strings."
 
-        corpus = self.request("RetrieveCorpus", id=self.corpus_id)
-        available_slugs = {element_type["slug"] for element_type in corpus["types"]}
-        missing_slugs = set(type_slugs) - available_slugs
+        if not self.corpus_types:
+            self.list_corpus_types()
 
+        missing_slugs = set(type_slugs) - set(self.corpus_types)
         if missing_slugs:
             if create_missing:
                 self.create_required_types(
@@ -79,7 +94,7 @@ class ElementMixin:
                 )
             else:
                 raise MissingTypeError(
-                    f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
+                    f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in corpus ({self.corpus_id}).'
                 )
 
         return True
@@ -145,7 +160,7 @@ class ElementMixin:
             logger.warning("Cannot create element as this worker is in read-only mode")
             return
 
-        sub_element = self.request(
+        sub_element = self.api_client.request(
             "CreateElement",
             body={
                 "type": type,
@@ -243,7 +258,7 @@ class ElementMixin:
             logger.warning("Cannot create elements as this worker is in read-only mode")
             return
 
-        created_ids = self.request(
+        created_ids = self.api_client.request(
             "CreateElements",
             id=parent.id,
             body={
@@ -311,7 +326,7 @@ class ElementMixin:
             logger.warning("Cannot link elements as this worker is in read-only mode")
             return
 
-        return self.request(
+        return self.api_client.request(
             "CreateElementParent",
             parent=parent.id,
             child=child.id,
@@ -383,7 +398,7 @@ class ElementMixin:
             logger.warning("Cannot update element as this worker is in read-only mode")
             return
 
-        updated_element = self.request(
+        updated_element = self.api_client.request(
             "PartialUpdateElement",
             id=element.id,
             body=kwargs,

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/entity.py

@@ -48,6 +48,7 @@ class EntityMixin:
         if not self.entity_types:
             # Load entity_types of corpus
             self.list_corpus_entity_types()
+
         for entity_type in entity_types:
             # Do nothing if type already exists
             if entity_type in self.entity_types:
@@ -60,7 +61,7 @@ class EntityMixin:
                 )
 
             # Create type if non-existent
-            self.entity_types[entity_type] = self.request(
+            self.entity_types[entity_type] = self.api_client.request(
                 "CreateEntityType",
                 body={
                     "name": entity_type,
@@ -106,7 +107,7 @@ class EntityMixin:
         entity_type_id = self.entity_types.get(type)
         assert entity_type_id, f"Entity type `{type}` not found in the corpus."
 
-        entity = self.request(
+        entity = self.api_client.request(
             "CreateEntity",
             body={
                 "name": name,
@@ -188,7 +189,7 @@ class EntityMixin:
         if confidence is not None:
             body["confidence"] = confidence
 
-        transcription_ent = self.request(
+        transcription_ent = self.api_client.request(
             "CreateTranscriptionEntity",
             id=transcription.id,
             body=body,
@@ -289,7 +290,7 @@ class EntityMixin:
             )
             return
 
-        created_ids = self.request(
+        created_ids = self.api_client.request(
             "CreateTranscriptionEntities",
             id=transcription.id,
             body={
@@ -385,9 +386,7 @@ class EntityMixin:
             f'Loaded {count} entit{"ies" if count > 1 else "y"} in corpus ({self.corpus_id})'
         )
 
-    def list_corpus_entity_types(
-        self,
-    ):
+    def list_corpus_entity_types(self):
         """
         Loads available entity types in corpus.
         """

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/metadata.py

@@ -93,7 +93,7 @@ class MetaDataMixin:
             logger.warning("Cannot create metadata as this worker is in read-only mode")
             return
 
-        metadata = self.request(
+        metadata = self.api_client.request(
             "CreateMetaData",
             id=element.id,
             body={
@@ -168,7 +168,7 @@ class MetaDataMixin:
             logger.warning("Cannot create metadata as this worker is in read-only mode")
             return
 
-        created_metadata_list = self.request(
+        created_metadata_list = self.api_client.request(
             "CreateMetaDataBulk",
             id=element.id,
             body={

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/task.py

@@ -22,7 +22,7 @@ class TaskMixin:
             task_id, uuid.UUID
         ), "task_id shouldn't be null and should be an UUID"
 
-        results = self.request("ListArtifacts", id=task_id)
+        results = self.api_client.request("ListArtifacts", id=task_id)
 
         return map(Artifact, results)
 
@@ -43,4 +43,6 @@ class TaskMixin:
             artifact, Artifact
         ), "artifact shouldn't be null and should be an Artifact"
 
-        return self.request("DownloadArtifact", id=task_id, path=artifact.path)
+        return self.api_client.request(
+            "DownloadArtifact", id=task_id, path=artifact.path
+        )

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/training.py

@@ -185,7 +185,7 @@ class TrainingMixin:
         assert not self.model_version, "A model version has already been created."
 
         configuration = configuration or {}
-        self.model_version = self.request(
+        self.model_version = self.api_client.request(
             "CreateModelVersion",
             id=model_id,
             body=build_clean_payload(
@@ -217,7 +217,7 @@ class TrainingMixin:
         :param parent: ID of the parent model version
         """
         assert self.model_version, "No model version has been created yet."
-        self.model_version = self.request(
+        self.model_version = self.api_client.request(
             "UpdateModelVersion",
             id=self.model_version["id"],
             body=build_clean_payload(
@@ -273,41 +273,44 @@ class TrainingMixin:
         """
         assert self.model_version, "You must create the model version and upload its archive before validating it."
         try:
-            self.model_version = self.request(
-                "ValidateModelVersion",
+            self.model_version = self.api_client.request(
+                "PartialUpdateModelVersion",
                 id=self.model_version["id"],
                 body={
+                    "state": "available",
                     "size": size,
                     "hash": hash,
                     "archive_hash": archive_hash,
                 },
             )
         except ErrorResponse as e:
-            # Temporary fix while waiting for `ValidateModelVersion` refactoring as it can
-            # return errors even when the model version is properly validated
-            if e.status_code in [403, 500]:
-                logger.warning(
-                    f'An error occurred while validating model version {self.model_version["id"]}, please check its status.'
-                )
-                return
-
-            if e.status_code != 409:
+            model_version = e.content
+            if not model_version or "id" not in model_version:
                 raise e
 
             logger.warning(
                 f"An available model version exists with hash {hash}, using it instead of the pending version."
             )
             pending_version_id = self.model_version["id"]
-            self.model_version = getattr(e, "content", None)
-            assert self.model_version is not None, "An unexpected error occurred."
-
             logger.warning("Removing the pending model version.")
             try:
-                self.request("DestroyModelVersion", id=pending_version_id)
+                self.api_client.request("DestroyModelVersion", id=pending_version_id)
             except ErrorResponse as e:
                 msg = getattr(e, "content", str(e))
                 logger.error(
                     f"An error occurred removing the pending version {pending_version_id}: {msg}."
                 )
 
+            logger.info("Retrieving the existing model version.")
+            existing_version_id = model_version["id"].pop()
+            try:
+                self.model_version = self.api_client.request(
+                    "RetrieveModelVersion", id=existing_version_id
+                )
+            except ErrorResponse as e:
+                logger.error(
+                    f"An error occurred retrieving the existing version {existing_version_id}: {e.status_code} - {e.content}."
+                )
+                raise
+
         logger.info(f"Model version {self.model_version['id']} is now available.")

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/transcription.py

@@ -77,7 +77,7 @@ class TranscriptionMixin:
             )
             return
 
-        created = self.request(
+        created = self.api_client.request(
             "CreateTranscription",
             id=element.id,
             body={
@@ -171,7 +171,7 @@ class TranscriptionMixin:
             )
             return
 
-        created_trs = self.request(
+        created_trs = self.api_client.request(
             "CreateTranscriptions",
             body={
                 "worker_run_id": self.worker_run_id,
@@ -291,7 +291,7 @@ class TranscriptionMixin:
             )
             return
 
-        annotations = self.request(
+        annotations = self.api_client.request(
             "CreateElementTranscriptions",
             id=element.id,
             body={

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/arkindex_worker/worker/version.py

@@ -34,7 +34,9 @@ class WorkerVersionMixin:
         if worker_version_id in self._worker_version_cache:
             return self._worker_version_cache[worker_version_id]
 
-        worker_version = self.request("RetrieveWorkerVersion", id=worker_version_id)
+        worker_version = self.api_client.request(
+            "RetrieveWorkerVersion", id=worker_version_id
+        )
         self._worker_version_cache[worker_version_id] = worker_version
 
         return worker_version

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/pyproject.toml

@@ -4,14 +4,14 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "arkindex-base-worker"
-version = "0.4.0a1"
+version = "0.4.0b1"
 description = "Base Worker to easily build Arkindex ML workflows"
 license = { file = "LICENSE" }
 dependencies = [
     "peewee~=3.17",
-    "Pillow==10.3.0",
+    "Pillow==10.4.0",
     "python-gnupg==0.5.2",
-    "shapely==2.0.3",
+    "shapely==2.0.5",
     "teklia-toolbox==0.1.5",
     "zstandard==0.22.0",
 ]
@@ -42,12 +42,12 @@ Authors = "https://teklia.com"
 
 [project.optional-dependencies]
 docs = [
-    "black==24.4.0",
-    "mkdocs-material==9.5.17",
-    "mkdocstrings-python==1.9.2",
+    "black==24.4.2",
+    "mkdocs-material==9.5.31",
+    "mkdocstrings-python==1.10.7",
 ]
 tests = [
-    "pytest==8.1.1",
+    "pytest==8.3.2",
     "pytest-mock==3.14.0",
     "pytest-responses==0.5.1",
 ]

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/tests/test_base_worker.py

@@ -658,7 +658,7 @@ def test_find_extras_directory_not_found(monkeypatch, extras_path, exists, error
 def test_find_parents_file_paths(responses, mock_base_worker_with_cache, tmp_path):
     responses.add(
         responses.GET,
-        "http://testserver/api/v1/task/my_task/from-agent/",
+        "http://testserver/api/v1/task/my_task/",
         status=200,
         json={"parents": ["first", "second", "third"]},
     )

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/tests/test_elements_worker/test_elements.py

@@ -22,6 +22,24 @@ from tests import CORPUS_ID
 from . import BASE_API_CALLS
 
 
+def test_list_corpus_types(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/",
+        json={
+            "id": CORPUS_ID,
+            "types": [{"slug": "folder"}, {"slug": "page"}],
+        },
+    )
+
+    mock_elements_worker.list_corpus_types()
+
+    assert mock_elements_worker.corpus_types == {
+        "folder": {"slug": "folder"},
+        "page": {"slug": "page"},
+    }
+
+
 def test_check_required_types_argument_types(mock_elements_worker):
     with pytest.raises(
         AssertionError, match="At least one element type slug is required."
@@ -32,17 +50,11 @@ def test_check_required_types_argument_types(mock_elements_worker):
         mock_elements_worker.check_required_types("lol", 42)
 
 
-def test_check_required_types(responses, mock_elements_worker):
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/corpus/{CORPUS_ID}/",
-        json={
-            "id": CORPUS_ID,
-            "name": "Some Corpus",
-            "types": [{"slug": "folder"}, {"slug": "page"}],
-        },
-    )
-    mock_elements_worker.setup_api_client()
+def test_check_required_types(mock_elements_worker):
+    mock_elements_worker.corpus_types = {
+        "folder": {"slug": "folder"},
+        "page": {"slug": "page"},
+    }
 
     assert mock_elements_worker.check_required_types("page")
     assert mock_elements_worker.check_required_types("page", "folder")
@@ -50,22 +62,18 @@ def test_check_required_types(responses, mock_elements_worker):
     with pytest.raises(
         MissingTypeError,
         match=re.escape(
-            "Element type(s) act, text_line were not found in the Some Corpus corpus (11111111-1111-1111-1111-111111111111)."
+            "Element type(s) act, text_line were not found in corpus (11111111-1111-1111-1111-111111111111)."
         ),
     ):
         assert mock_elements_worker.check_required_types("page", "text_line", "act")
 
 
 def test_create_missing_types(responses, mock_elements_worker):
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/corpus/{CORPUS_ID}/",
-        json={
-            "id": CORPUS_ID,
-            "name": "Some Corpus",
-            "types": [{"slug": "folder"}, {"slug": "page"}],
-        },
-    )
+    mock_elements_worker.corpus_types = {
+        "folder": {"slug": "folder"},
+        "page": {"slug": "page"},
+    }
+
     responses.add(
         responses.POST,
         "http://testserver/api/v1/elements/type/",
@@ -94,7 +102,6 @@ def test_create_missing_types(responses, mock_elements_worker):
             )
         ],
     )
-    mock_elements_worker.setup_api_client()
 
     assert mock_elements_worker.check_required_types(
         "page", "text_line", "act", create_missing=True

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/tests/test_elements_worker/test_training.py

@@ -179,44 +179,12 @@ def test_validate_model_version_not_created(mock_training_worker):
         mock_training_worker.validate_model_version(hash="a", size=1, archive_hash="b")
 
 
-@pytest.mark.parametrize("status_code", [403, 500])
-def test_validate_model_version_catch_errors(
-    mocker, mock_training_worker, caplog, status_code
-):
-    mocker.patch(
-        "arkindex.client.ArkindexClient.request.retry.retry", return_value=False
-    )
-
-    mock_training_worker.model_version = {"id": "model_version_id"}
-    args = {
-        "hash": "hash",
-        "archive_hash": "archive_hash",
-        "size": 30,
-    }
-    mock_training_worker.api_client.add_error_response(
-        "ValidateModelVersion",
-        id="model_version_id",
-        status_code=status_code,
-        body=args,
-    )
-
-    mock_training_worker.validate_model_version(**args)
-    assert mock_training_worker.model_version == {"id": "model_version_id"}
-    assert [
-        (level, message)
-        for module, level, message in caplog.record_tuples
-        if module == "arkindex_worker"
-    ] == [
-        (
-            logging.WARNING,
-            "An error occurred while validating model version model_version_id, please check its status.",
-        ),
-    ]
-
-
 @pytest.mark.parametrize("deletion_failed", [True, False])
 def test_validate_model_version_hash_conflict(
-    mock_training_worker, default_model_version, caplog, deletion_failed
+    mock_training_worker,
+    default_model_version,
+    caplog,
+    deletion_failed,
 ):
     mock_training_worker.model_version = {"id": "another_id"}
     args = {
@@ -225,11 +193,11 @@ def test_validate_model_version_hash_conflict(
         "size": 30,
     }
     mock_training_worker.api_client.add_error_response(
-        "ValidateModelVersion",
+        "PartialUpdateModelVersion",
         id="another_id",
         status_code=409,
-        body=args,
-        content=default_model_version,
+        body={"state": "available", **args},
+        content={"id": ["model_version_id"]},
     )
     if deletion_failed:
         mock_training_worker.api_client.add_error_response(
@@ -244,6 +212,11 @@ def test_validate_model_version_hash_conflict(
             id="another_id",
             response="No content",
         )
+    mock_training_worker.api_client.add_response(
+        "RetrieveModelVersion",
+        id="model_version_id",
+        response=default_model_version,
+    )
 
     mock_training_worker.validate_model_version(**args)
     assert mock_training_worker.model_version == default_model_version
@@ -266,6 +239,7 @@ def test_validate_model_version_hash_conflict(
         ),
         (logging.WARNING, "Removing the pending model version."),
         *error_msg,
+        (logging.INFO, "Retrieving the existing model version."),
         (logging.INFO, "Model version model_version_id is now available."),
     ]
 
@@ -278,9 +252,9 @@ def test_validate_model_version(mock_training_worker, default_model_version, cap
         "size": 30,
     }
     mock_training_worker.api_client.add_response(
-        "ValidateModelVersion",
+        "PartialUpdateModelVersion",
         id="model_version_id",
-        body=args,
+        body={"state": "available", **args},
         response=default_model_version,
     )
 

{arkindex_base_worker-0.4.0a1 → arkindex_base_worker-0.4.0b1}/tests/test_merge.py

@@ -161,7 +161,7 @@ def test_merge_from_worker(
     """
     responses.add(
         responses.GET,
-        "http://testserver/api/v1/task/my_task/from-agent/",
+        "http://testserver/api/v1/task/my_task/",
         status=200,
         json={"parents": ["first", "second"]},
     )