arkindex-base-worker 0.4.0b1__py3-none-any.whl → 0.4.0b2__py3-none-any.whl

tests/test_elements_worker/test_entities.py

@@ -13,6 +13,7 @@ from arkindex_worker.cache import (
     CachedTranscriptionEntity,
 )
 from arkindex_worker.models import Transcription
+from arkindex_worker.utils import DEFAULT_BATCH_SIZE
 from arkindex_worker.worker.entity import MissingEntityType
 from arkindex_worker.worker.transcription import TextOrientation
 from tests import CORPUS_ID
@@ -988,38 +989,89 @@ def test_create_transcription_entities_wrong_entity(
         )
 
 
-def test_create_transcription_entities(responses, mock_elements_worker):
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+def test_create_transcription_entities(batch_size, responses, mock_elements_worker):
     transcription = Transcription(id="transcription-id")
+
     # Call to Transcription entities creation in bulk
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/transcription/transcription-id/entities/bulk/",
-        status=201,
-        match=[
-            matchers.json_params_matcher(
-                {
-                    "worker_run_id": "56785678-5678-5678-5678-567856785678",
+    if batch_size > 1:
+        responses.add(
+            responses.POST,
+            "http://testserver/api/v1/transcription/transcription-id/entities/bulk/",
+            status=201,
+            match=[
+                matchers.json_params_matcher(
+                    {
+                        "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                        "entities": [
+                            {
+                                "name": "Teklia",
+                                "type_id": "22222222-2222-2222-2222-222222222222",
+                                "offset": 0,
+                                "length": 6,
+                                "confidence": 1.0,
+                            },
+                            {
+                                "name": "Team Rocket",
+                                "type_id": "22222222-2222-2222-2222-222222222222",
+                                "offset": 7,
+                                "length": 11,
+                                "confidence": 1.0,
+                            },
+                        ],
+                    }
+                )
+            ],
+            json={
+                "entities": [
+                    {
+                        "transcription_entity_id": "transc-entity-id",
+                        "entity_id": "entity-id1",
+                    },
+                    {
+                        "transcription_entity_id": "transc-entity-id",
+                        "entity_id": "entity-id2",
+                    },
+                ]
+            },
+        )
+    else:
+        for idx, (name, offset, length) in enumerate(
+            [
+                ("Teklia", 0, 6),
+                ("Team Rocket", 7, 11),
+            ],
+            start=1,
+        ):
+            responses.add(
+                responses.POST,
+                "http://testserver/api/v1/transcription/transcription-id/entities/bulk/",
+                status=201,
+                match=[
+                    matchers.json_params_matcher(
+                        {
+                            "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                            "entities": [
+                                {
+                                    "name": name,
+                                    "type_id": "22222222-2222-2222-2222-222222222222",
+                                    "offset": offset,
+                                    "length": length,
+                                    "confidence": 1.0,
+                                }
+                            ],
+                        }
+                    )
+                ],
+                json={
                     "entities": [
                         {
-                            "name": "Teklia",
-                            "type_id": "22222222-2222-2222-2222-222222222222",
-                            "offset": 0,
-                            "length": 6,
-                            "confidence": 1.0,
+                            "transcription_entity_id": "transc-entity-id",
+                            "entity_id": f"entity-id{idx}",
                         }
-                    ],
-                }
+                    ]
+                },
             )
-        ],
-        json={
-            "entities": [
-                {
-                    "transcription_entity_id": "transc-entity-id",
-                    "entity_id": "entity-id",
-                }
-            ]
-        },
-    )
 
     # Store entity type/slug correspondence on the worker
     mock_elements_worker.entity_types = {
@@ -1034,18 +1086,35 @@ def test_create_transcription_entities(responses, mock_elements_worker):
                 "offset": 0,
                 "length": 6,
                 "confidence": 1.0,
-            }
+            },
+            {
+                "name": "Team Rocket",
+                "type_id": "22222222-2222-2222-2222-222222222222",
+                "offset": 7,
+                "length": 11,
+                "confidence": 1.0,
+            },
         ],
+        batch_size=batch_size,
     )
 
-    assert len(created_objects) == 1
+    assert len(created_objects) == 2
 
-    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/transcription/transcription-id/entities/bulk/",
-        ),
+        )
     ]
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(
+            (
+                "POST",
+                "http://testserver/api/v1/transcription/transcription-id/entities/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

tests/test_elements_worker/test_metadata.py

@@ -7,6 +7,7 @@ from apistar.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
@@ -228,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"],
+        },
     ]