arkindex-base-worker 0.3.7rc4__py3-none-any.whl → 0.5.0a1__py3-none-any.whl

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"]},
     )
@@ -181,7 +181,6 @@ def test_merge_from_worker(
     (tmp_path / "my_task").mkdir()
     mock_base_worker_with_cache.args = mock_base_worker_with_cache.parser.parse_args()
     mock_base_worker_with_cache.configure()
-    mock_base_worker_with_cache.configure_cache()
     # Store parent tasks IDs as attribute
     assert mock_base_worker_with_cache.task_parents == ["first", "second"]
 

tests/test_utils.py

@@ -1,9 +1,34 @@
-from pathlib import Path
+import logging
 
-from arkindex_worker.utils import close_delete_file, extract_tar_zst_archive
+import pytest
 
-FIXTURES = Path(__file__).absolute().parent / "data"
-ARCHIVE = FIXTURES / "archive.tar.zst"
+from arkindex_worker.cache import unsupported_cache
+from arkindex_worker.utils import (
+    DEFAULT_BATCH_SIZE,
+    batch_publication,
+    close_delete_file,
+    extract_tar_zst_archive,
+    parse_source_id,
+)
+from tests import FIXTURES_DIR
+
+ARCHIVE = FIXTURES_DIR / "archive.tar.zst"
+
+
+@pytest.mark.parametrize(
+    ("source_id", "expected"),
+    [
+        (None, None),
+        ("", None),
+        (
+            "cafecafe-cafe-cafe-cafe-cafecafecafe",
+            "cafecafe-cafe-cafe-cafe-cafecafecafe",
+        ),
+        ("manual", False),
+    ],
+)
+def test_parse_source_id(source_id, expected):
+    assert parse_source_id(source_id) == expected
 
 
 def test_extract_tar_zst_archive(tmp_path):
@@ -33,3 +58,63 @@ def test_close_delete_file(tmp_path):
     close_delete_file(archive_fd, archive_path)
 
     assert not archive_path.exists()
+
+
+class TestMixin:
+    def __init__(self, use_cache: bool = False):
+        """
+        Args:
+            use_cache (bool, optional): To mock BaseWorker.use_cache attribute. Defaults to False.
+        """
+        self.use_cache = use_cache
+
+    @batch_publication
+    def custom_publication_in_batches(self, batch_size: int = DEFAULT_BATCH_SIZE):
+        return batch_size
+
+    @unsupported_cache
+    @batch_publication
+    def custom_publication_in_batches_without_cache(
+        self, batch_size: int = DEFAULT_BATCH_SIZE
+    ):
+        return batch_size
+
+
+def test_batch_publication_decorator_no_parameter():
+    assert TestMixin().custom_publication_in_batches() == DEFAULT_BATCH_SIZE
+
+
+@pytest.mark.parametrize("wrong_batch_size", [None, "not an int", 0])
+def test_batch_publication_decorator_wrong_parameter(wrong_batch_size):
+    with pytest.raises(
+        AssertionError,
+        match="batch_size shouldn't be null and should be a strictly positive integer",
+    ):
+        TestMixin().custom_publication_in_batches(batch_size=wrong_batch_size)
+
+
+@pytest.mark.parametrize("batch_size", [1, 10, DEFAULT_BATCH_SIZE])
+def test_batch_publication_decorator_right_parameter(batch_size):
+    assert (
+        TestMixin().custom_publication_in_batches(batch_size=batch_size) == batch_size
+    )
+
+
+def test_batch_publication_decorator_alongside_unsupported_cache(caplog):
+    # Capture log messages
+    caplog.clear()
+    with caplog.at_level(logging.WARNING):
+        # Call the helper
+        assert (
+            TestMixin(use_cache=True).custom_publication_in_batches_without_cache()
+            == DEFAULT_BATCH_SIZE
+        )
+
+    # Check logs
+    assert caplog.record_tuples == [
+        (
+            "arkindex_worker",
+            logging.WARNING,
+            "This API helper `custom_publication_in_batches_without_cache` did not update the cache database",
+        ),
+    ]

worker-demo/tests/conftest.py

@@ -0,0 +1,32 @@
+import os
+
+import pytest
+
+from arkindex.mock import MockApiClient
+from arkindex_worker.worker.base import BaseWorker
+
+
+@pytest.fixture(autouse=True)
+def _setup_environment(responses, monkeypatch) -> None:
+    """Setup needed environment variables"""
+
+    # Allow accessing remote API schemas
+    # defaulting to the prod environment
+    schema_url = os.environ.get(
+        "ARKINDEX_API_SCHEMA_URL",
+        "https://demo.arkindex.org/api/v1/openapi/?format=openapi-json",
+    )
+    responses.add_passthru(schema_url)
+
+    # Set schema url in environment
+    os.environ["ARKINDEX_API_SCHEMA_URL"] = schema_url
+    # Setup a fake worker run ID
+    os.environ["ARKINDEX_WORKER_RUN_ID"] = "1234-demo"
+    # Setup a fake corpus ID
+    os.environ["ARKINDEX_CORPUS_ID"] = "1234-corpus-id"
+
+    # Setup a mock api client instead of using a real one
+    def mock_setup_api_client(self):
+        self.api_client = MockApiClient()
+
+    monkeypatch.setattr(BaseWorker, "setup_api_client", mock_setup_api_client)

worker-demo/tests/test_worker.py

@@ -0,0 +1,12 @@
+import importlib
+
+
+def test_dummy():
+    assert True
+
+
+def test_import():
+    """Import our newly created module, through importlib to avoid parsing issues"""
+    worker = importlib.import_module("worker_demo.worker")
+    assert hasattr(worker, "Demo")
+    assert hasattr(worker.Demo, "process_element")

worker-demo/worker_demo/__init__.py

@@ -0,0 +1,6 @@
+import logging
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s/%(name)s: %(message)s",
+)

worker-demo/worker_demo/worker.py

@@ -0,0 +1,19 @@
+from logging import Logger, getLogger
+
+from arkindex_worker.models import Element
+from arkindex_worker.worker import ElementsWorker
+
+logger: Logger = getLogger(__name__)
+
+
+class Demo(ElementsWorker):
+    def process_element(self, element: Element) -> None:
+        logger.info(f"Demo processing element ({element.id})")
+
+
+def main() -> None:
+    Demo(description="Demo ML worker for Arkindex").run()
+
+
+if __name__ == "__main__":
+    main()

arkindex_base_worker-0.3.7rc4.dist-info/RECORD

@@ -1,41 +0,0 @@
-arkindex_worker/__init__.py,sha256=OlgCtTC9MaWeejviY0a3iQpALcRQGMVArFVVYwTF6I8,162
-arkindex_worker/cache.py,sha256=FTlB0coXofn5zTNRTcVIvh709mcw4a1bPGqkwWjKs3w,11248
-arkindex_worker/image.py,sha256=9-k_Wojk-sLbgvBSi7tWiiDc9YApWauJpHGKRay_nmo,14166
-arkindex_worker/models.py,sha256=HdKFw3qk4WIWC-DrHDkhsw0mHP3OILuCLFf7aTjruZU,9526
-arkindex_worker/utils.py,sha256=VSO8c21nsSaUCkyJaFX8wOwDQ0tztLOBFtiGvqlT0zU,6900
-arkindex_worker/worker/__init__.py,sha256=I8QmdAs659SalxNjtCu2K2ItdyUlXYm3mK_WhZdjgBs,19498
-arkindex_worker/worker/base.py,sha256=7ii3rZai6IB0-eB0TJ6pg-IhxMmW4izoJAKJKczbyZ4,19934
-arkindex_worker/worker/classification.py,sha256=0OiwxV9lb97Zs3kODm3hzyk0V7IxBTiW5SL6AYgRH1M,10351
-arkindex_worker/worker/dataset.py,sha256=qzjaXJtfeNCP2acsHbqp5tjQk-KpLHwVzjDAExeAmVg,3228
-arkindex_worker/worker/element.py,sha256=AWK3YJSHWy3j4ajntJloi_2X4zxsgXZ6c6dzphgq3OI,33848
-arkindex_worker/worker/entity.py,sha256=YT2Ttdn-L5TRoDdhOI3Z4GE1vtkWl7tKZqbYrtxZ2Ug,14630
-arkindex_worker/worker/metadata.py,sha256=SC6apVaOjFrmYw5b-njhqIlH-_r0ExbNpZeQZzlUjBE,6669
-arkindex_worker/worker/task.py,sha256=cz3wJNPgogZv1lm_3lm7WScitQtYQtL6H6I7Xokq208,1475
-arkindex_worker/worker/training.py,sha256=SOs3YKGikTr3rdWYp9H-jbtgRnZxQAoqtwB26ztx9j8,10235
-arkindex_worker/worker/transcription.py,sha256=6R7ofcGnNqX4rjT0kRKIE-G9FHq2TJ1tfztNM5sTqYE,20464
-arkindex_worker/worker/version.py,sha256=cs2pdlDxpKRO2Oldvcu54w-D_DQhf1cdeEt4tKX_QYs,1927
-tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/conftest.py,sha256=wzKXRnS7OKQCNHrlDaQhMC8EXlsQTY_S4L9U_hXbjpM,22004
-tests/test_base_worker.py,sha256=Uq6_MpLW23gmKFXkU-SyDUaA_4dlViLBGG4e3gpBBz0,24512
-tests/test_cache.py,sha256=ii0gyr0DrG7ChEs7pmT8hMdSguAOAcCze4bRMiFQxuk,10640
-tests/test_dataset_worker.py,sha256=Q-3gVu2FNa6mJVkUW-PUVgyUAvRkxSzLCJrPhwoJlxQ,28273
-tests/test_element.py,sha256=2G9M15TLxQRmvrWM9Kw2ucnElh4kSv_oF_5FYwwAxTY,13181
-tests/test_image.py,sha256=FZv8njLxh45sVgmY71UFHt0lv1cHr0cK4rrtPhQleX8,16262
-tests/test_merge.py,sha256=Q4zCbtZbe0wBfqE56gvAD06c6pDuhqnjKaioFqIgAQw,8331
-tests/test_utils.py,sha256=pFXegcBvIuy1tJDDSgQtCbC_tRaoLjd2055R5lu3hS0,1236
-tests/test_elements_worker/__init__.py,sha256=Fh4nkbbyJSMv_VtjQxnWrOqTnxXaaWI8S9WU0VrzCHs,179
-tests/test_elements_worker/test_classifications.py,sha256=PE88fsdra8QsWcKjSyao-pTHlaIWNxlbfF0CrLe9LBA,26517
-tests/test_elements_worker/test_cli.py,sha256=BsFTswLti63WAZ2pf6ipiZKWJJyCQuSfuKnSlESuK8g,2878
-tests/test_elements_worker/test_dataset.py,sha256=-kVll1NcMPWkIx8D7r-Z5neEGkFiZ9YQfC4eTMIfjg0,13475
-tests/test_elements_worker/test_elements.py,sha256=6XKtgXSVQJnTSgTHWwEVsAtIwLBapjYjUYPUdjxcHsY,84971
-tests/test_elements_worker/test_entities.py,sha256=yi1mXzvKvNwUNMzo0UZ56YOIJstYHcLyeepPJ8f10MQ,34557
-tests/test_elements_worker/test_metadata.py,sha256=b9CNv4W31TRJqYauvX_pRIN2SvnybaLqF-FWoFwa2Vc,18672
-tests/test_elements_worker/test_task.py,sha256=FCpxE9UpouKXgjGvWgNHEai_Hiy2d1YmqRG-_v2s27s,6312
-tests/test_elements_worker/test_training.py,sha256=WeG-cDuJ-YhPgfKH47TtXBxyargtLuk7c8tsik2WnL8,8414
-tests/test_elements_worker/test_transcriptions.py,sha256=WVJG26sZyY66fu-Eka9A1_WWIeNI2scogjypzURnp8A,73468
-tests/test_elements_worker/test_worker.py,sha256=7-jGJVT3yMGpIyN96Uafz5eIUrO4ieNLgw0k1D8BhGc,17163
-arkindex_base_worker-0.3.7rc4.dist-info/LICENSE,sha256=NVshRi1efwVezMfW7xXYLrdDr2Li1AfwfGOd5WuH1kQ,1063
-arkindex_base_worker-0.3.7rc4.dist-info/METADATA,sha256=ilh4IdFYSXepgr0imEMH3ZbewlFJlbg97VKnvhKXMVQ,3411
-arkindex_base_worker-0.3.7rc4.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-arkindex_base_worker-0.3.7rc4.dist-info/top_level.txt,sha256=TtagLI8LSv7GE7nG8MQqDFAJ5bNDPJn7Z5vizOgrWkA,22
-arkindex_base_worker-0.3.7rc4.dist-info/RECORD,,