arkindex-base-worker 0.3.7rc9__py3-none-any.whl → 0.4.0__py3-none-any.whl

tests/test_elements_worker/test_metadata.py

@@ -2,11 +2,12 @@ import json
 import re
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
+from arkindex.exceptions import ErrorResponse
 from arkindex.mock import MockApiClient
 from arkindex_worker.cache import CachedElement
 from arkindex_worker.models import Element
+from arkindex_worker.utils import DEFAULT_BATCH_SIZE
 from arkindex_worker.worker import MetaType
 
 from . import BASE_API_CALLS
@@ -136,7 +137,7 @@ def test_create_metadata_api_error(responses, mock_elements_worker):
     responses.add(
         responses.POST,
         "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
@@ -147,31 +148,14 @@ def test_create_metadata_api_error(responses, mock_elements_worker):
             value="La Turbine, Grenoble 38000",
         )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        # We retry 5 times the API call
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/",
-        ),
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/",
-        ),
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/",
-        ),
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/",
-        ),
         (
             "POST",
             "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/",
-        ),
+        )
     ]
 
 
@@ -245,143 +229,267 @@ def test_create_metadata_cached_element(responses, mock_elements_worker_with_cac
     assert metadata_id == "12345678-1234-1234-1234-123456789123"
 
 
-@pytest.mark.parametrize(
-    "metadata_list",
-    [
-        [{"type": MetaType.Text, "name": "fake_name", "value": "fake_value"}],
-        [
-            {
-                "type": MetaType.Text,
-                "name": "fake_name",
-                "value": "fake_value",
-                "entity_id": "fake_entity_id",
-            }
-        ],
-    ],
-)
-def test_create_metadata_bulk(responses, mock_elements_worker, metadata_list):
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+def test_create_metadata_bulk(batch_size, responses, mock_elements_worker):
     element = Element({"id": "12341234-1234-1234-1234-123412341234"})
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        status=201,
-        json={
-            "worker_run_id": "56785678-5678-5678-5678-567856785678",
-            "metadata_list": [
-                {
-                    "id": "fake_metadata_id",
-                    "type": metadata_list[0]["type"].value,
-                    "name": metadata_list[0]["name"],
-                    "value": metadata_list[0]["value"],
-                    "dates": [],
-                    "entity_id": metadata_list[0].get("entity_id"),
-                }
-            ],
+
+    metadata_list = [
+        {"type": MetaType.Text, "name": "fake_name", "value": "fake_value"},
+        {
+            "type": MetaType.Text,
+            "name": "Year",
+            "value": "2024",
+            "entity_id": "entity_id",
         },
-    )
+    ]
+    if batch_size > 1:
+        responses.add(
+            responses.POST,
+            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
+            status=201,
+            json={
+                "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                "metadata_list": [
+                    {
+                        "id": "fake_metadata_id1",
+                        "type": metadata_list[0]["type"].value,
+                        "name": metadata_list[0]["name"],
+                        "value": metadata_list[0]["value"],
+                        "dates": [],
+                        "entity_id": None,
+                    },
+                    {
+                        "id": "fake_metadata_id2",
+                        "type": metadata_list[1]["type"].value,
+                        "name": metadata_list[1]["name"],
+                        "value": metadata_list[1]["value"],
+                        "dates": [],
+                        "entity_id": metadata_list[1]["entity_id"],
+                    },
+                ],
+            },
+        )
+    else:
+        for idx, meta in enumerate(metadata_list, start=1):
+            responses.add(
+                responses.POST,
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
+                status=201,
+                json={
+                    "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                    "metadata_list": [
+                        {
+                            "id": f"fake_metadata_id{idx}",
+                            "type": meta["type"].value,
+                            "name": meta["name"],
+                            "value": meta["value"],
+                            "dates": [],
+                            "entity_id": meta.get("entity_id"),
+                        }
+                    ],
+                },
+            )
 
     created_metadata_list = mock_elements_worker.create_metadata_bulk(
-        element, metadata_list
+        element, metadata_list, batch_size=batch_size
     )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 1
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
+    bulk_api_calls = [
         (
             "POST",
             "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
-    ]
-    assert json.loads(responses.calls[-1].request.body)["metadata_list"] == [
-        {
-            "type": metadata_list[0]["type"].value,
-            "name": metadata_list[0]["name"],
-            "value": metadata_list[0]["value"],
-            "entity_id": metadata_list[0].get("entity_id"),
-        }
+        )
     ]
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(
+            (
+                "POST",
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
+            )
+        )
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + len(bulk_api_calls)
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + bulk_api_calls
+
+    first_meta = {
+        **metadata_list[0],
+        "type": metadata_list[0]["type"].value,
+        "entity_id": None,
+    }
+    second_meta = {**metadata_list[1], "type": metadata_list[1]["type"].value}
+    empty_payload = {
+        "worker_run_id": "56785678-5678-5678-5678-567856785678",
+        "metadata_list": [],
+    }
+
+    bodies = []
+    first_call_idx = None
+    if batch_size > 1:
+        first_call_idx = -1
+        bodies.append({**empty_payload, "metadata_list": [first_meta, second_meta]})
+    else:
+        first_call_idx = -2
+        bodies.append({**empty_payload, "metadata_list": [first_meta]})
+        bodies.append({**empty_payload, "metadata_list": [second_meta]})
+
+    assert [
+        json.loads(bulk_call.request.body)
+        for bulk_call in responses.calls[first_call_idx:]
+    ] == bodies
+
     assert created_metadata_list == [
         {
-            "id": "fake_metadata_id",
+            "id": "fake_metadata_id1",
             "type": metadata_list[0]["type"].value,
             "name": metadata_list[0]["name"],
             "value": metadata_list[0]["value"],
             "dates": [],
-            "entity_id": metadata_list[0].get("entity_id"),
-        }
+            "entity_id": None,
+        },
+        {
+            "id": "fake_metadata_id2",
+            "type": metadata_list[1]["type"].value,
+            "name": metadata_list[1]["name"],
+            "value": metadata_list[1]["value"],
+            "dates": [],
+            "entity_id": metadata_list[1]["entity_id"],
+        },
     ]
 
 
-@pytest.mark.parametrize(
-    "metadata_list",
-    [
-        [{"type": MetaType.Text, "name": "fake_name", "value": "fake_value"}],
-        [
-            {
-                "type": MetaType.Text,
-                "name": "fake_name",
-                "value": "fake_value",
-                "entity_id": "fake_entity_id",
-            }
-        ],
-    ],
-)
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
 def test_create_metadata_bulk_cached_element(
-    responses, mock_elements_worker_with_cache, metadata_list
+    batch_size, responses, mock_elements_worker_with_cache
 ):
     element = CachedElement.create(
         id="12341234-1234-1234-1234-123412341234", type="thing"
     )
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        status=201,
-        json={
-            "worker_run_id": "56785678-5678-5678-5678-567856785678",
-            "metadata_list": [
-                {
-                    "id": "fake_metadata_id",
-                    "type": metadata_list[0]["type"].value,
-                    "name": metadata_list[0]["name"],
-                    "value": metadata_list[0]["value"],
-                    "dates": [],
-                    "entity_id": metadata_list[0].get("entity_id"),
-                }
-            ],
+
+    metadata_list = [
+        {"type": MetaType.Text, "name": "fake_name", "value": "fake_value"},
+        {
+            "type": MetaType.Text,
+            "name": "Year",
+            "value": "2024",
+            "entity_id": "entity_id",
         },
-    )
+    ]
+    if batch_size > 1:
+        responses.add(
+            responses.POST,
+            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
+            status=201,
+            json={
+                "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                "metadata_list": [
+                    {
+                        "id": "fake_metadata_id1",
+                        "type": metadata_list[0]["type"].value,
+                        "name": metadata_list[0]["name"],
+                        "value": metadata_list[0]["value"],
+                        "dates": [],
+                        "entity_id": None,
+                    },
+                    {
+                        "id": "fake_metadata_id2",
+                        "type": metadata_list[1]["type"].value,
+                        "name": metadata_list[1]["name"],
+                        "value": metadata_list[1]["value"],
+                        "dates": [],
+                        "entity_id": metadata_list[1]["entity_id"],
+                    },
+                ],
+            },
+        )
+    else:
+        for idx, meta in enumerate(metadata_list, start=1):
+            responses.add(
+                responses.POST,
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
+                status=201,
+                json={
+                    "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                    "metadata_list": [
+                        {
+                            "id": f"fake_metadata_id{idx}",
+                            "type": meta["type"].value,
+                            "name": meta["name"],
+                            "value": meta["value"],
+                            "dates": [],
+                            "entity_id": meta.get("entity_id"),
+                        }
+                    ],
+                },
+            )
 
     created_metadata_list = mock_elements_worker_with_cache.create_metadata_bulk(
-        element, metadata_list
+        element, metadata_list, batch_size=batch_size
     )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 1
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
+    bulk_api_calls = [
         (
             "POST",
             "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
-    ]
-    assert json.loads(responses.calls[-1].request.body)["metadata_list"] == [
-        {
-            "type": metadata_list[0]["type"].value,
-            "name": metadata_list[0]["name"],
-            "value": metadata_list[0]["value"],
-            "entity_id": metadata_list[0].get("entity_id"),
-        }
+        )
     ]
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(
+            (
+                "POST",
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
+            )
+        )
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + len(bulk_api_calls)
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + bulk_api_calls
+
+    first_meta = {
+        **metadata_list[0],
+        "type": metadata_list[0]["type"].value,
+        "entity_id": None,
+    }
+    second_meta = {**metadata_list[1], "type": metadata_list[1]["type"].value}
+    empty_payload = {
+        "worker_run_id": "56785678-5678-5678-5678-567856785678",
+        "metadata_list": [],
+    }
+
+    bodies = []
+    first_call_idx = None
+    if batch_size > 1:
+        first_call_idx = -1
+        bodies.append({**empty_payload, "metadata_list": [first_meta, second_meta]})
+    else:
+        first_call_idx = -2
+        bodies.append({**empty_payload, "metadata_list": [first_meta]})
+        bodies.append({**empty_payload, "metadata_list": [second_meta]})
+
+    assert [
+        json.loads(bulk_call.request.body)
+        for bulk_call in responses.calls[first_call_idx:]
+    ] == bodies
+
     assert created_metadata_list == [
         {
-            "id": "fake_metadata_id",
+            "id": "fake_metadata_id1",
             "type": metadata_list[0]["type"].value,
             "name": metadata_list[0]["name"],
             "value": metadata_list[0]["value"],
             "dates": [],
-            "entity_id": metadata_list[0].get("entity_id"),
-        }
+            "entity_id": None,
+        },
+        {
+            "id": "fake_metadata_id2",
+            "type": metadata_list[1]["type"].value,
+            "name": metadata_list[1]["name"],
+            "value": metadata_list[1]["value"],
+            "dates": [],
+            "entity_id": metadata_list[1]["entity_id"],
+        },
     ]
 
 
@@ -474,37 +582,20 @@ def test_create_metadata_bulk_api_error(responses, mock_elements_worker):
     responses.add(
         responses.POST,
         "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
         mock_elements_worker.create_metadata_bulk(element, metadata_list)
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        # We retry 5 times the API call
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
         (
             "POST",
             "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
-        (
-            "POST",
-            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        ),
+        )
     ]
 
 

tests/test_elements_worker/test_process.py

@@ -0,0 +1,89 @@
+import pytest
+
+from tests import PROCESS_ID
+
+
+@pytest.mark.parametrize(
+    ("with_image", "elements"),
+    [
+        (
+            False,
+            [
+                {
+                    "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "type_id": "baaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "name": "element 1",
+                    "confidence": 1,
+                    "image_id": None,
+                    "image_width": None,
+                    "image_height": None,
+                    "image_url": None,
+                    "polygon": None,
+                    "rotation_angle": 0,
+                    "mirrored": False,
+                },
+                {
+                    "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaa0",
+                    "type_id": "baaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "name": "element 2",
+                    "confidence": 1,
+                    "image_id": None,
+                    "image_width": None,
+                    "image_height": None,
+                    "image_url": None,
+                    "polygon": None,
+                    "rotation_angle": 0,
+                    "mirrored": False,
+                },
+            ],
+        ),
+        (
+            True,
+            [
+                {
+                    "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "type_id": "baaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "name": "element 1",
+                    "confidence": 1,
+                    "image_id": "aaa2aaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "image_width": 76,
+                    "image_height": 138,
+                    "image_url": "http://somewhere.com/iiif/image.jpeg",
+                    "polygon": [[0, 0], [0, 40], [20, 40], [20, 0]],
+                    "rotation_angle": 0,
+                    "mirrored": False,
+                },
+                {
+                    "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaa0",
+                    "type_id": "baaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "name": "element 2",
+                    "confidence": 1,
+                    "image_id": "aaa2aaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
+                    "image_width": 138,
+                    "image_height": 76,
+                    "image_url": "http://somewhere.com/iiif/image.jpeg",
+                    "polygon": [[0, 0], [0, 40], [20, 40], [20, 0]],
+                    "rotation_angle": 0,
+                    "mirrored": False,
+                },
+            ],
+        ),
+    ],
+)
+def test_list_process_elements_with_image(
+    responses, mock_elements_worker, with_image, elements
+):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/process/{PROCESS_ID}/elements/?page_size=500&with_count=true&with_image={with_image}",
+        status=200,
+        json={
+            "count": 2,
+            "next": None,
+            "results": elements,
+        },
+    )
+    assert (
+        list(mock_elements_worker.list_process_elements(with_image=with_image))
+        == elements
+    )

tests/test_elements_worker/test_task.py

@@ -1,10 +1,10 @@
 import uuid
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
+from arkindex.exceptions import ErrorResponse
 from arkindex_worker.models import Artifact
-from tests.conftest import FIXTURES_DIR
+from tests import FIXTURES_DIR
 from tests.test_elements_worker import BASE_API_CALLS
 
 TASK_ID = uuid.UUID("cafecafe-cafe-cafe-cafe-cafecafecafe")
@@ -33,22 +33,17 @@ def test_list_artifacts_api_error(responses, mock_dataset_worker):
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/task/{TASK_ID}/artifacts/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
         mock_dataset_worker.list_artifacts(task_id=TASK_ID)
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifacts/"),
+        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifacts/")
     ]
 
 
@@ -153,7 +148,7 @@ def test_download_artifact_api_error(responses, mock_dataset_worker, default_art
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
@@ -161,16 +156,11 @@ def test_download_artifact_api_error(responses, mock_dataset_worker, default_art
             task_id=TASK_ID, artifact=default_artifact
         )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst"),
+        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst")
     ]