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

tests/test_elements_worker/test_elements.py

@@ -15,6 +15,7 @@ from arkindex_worker.cache import (
     init_cache_db,
 )
 from arkindex_worker.models import Element
+from arkindex_worker.utils import DEFAULT_BATCH_SIZE
 from arkindex_worker.worker import ElementsWorker
 from arkindex_worker.worker.element import MissingTypeError
 from tests import CORPUS_ID
@@ -22,6 +23,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 +51,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 +63,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 types 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 +103,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
@@ -1003,7 +1011,10 @@ def test_create_elements_api_error(responses, mock_elements_worker):
     ]
 
 
-def test_create_elements_cached_element(responses, mock_elements_worker_with_cache):
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+def test_create_elements_cached_element(
+    batch_size, responses, mock_elements_worker_with_cache
+):
     image = CachedImage.create(
         id=UUID("c0fec0fe-c0fe-c0fe-c0fe-c0fec0fec0fe"),
         width=42,
@@ -1016,12 +1027,28 @@ def test_create_elements_cached_element(responses, mock_elements_worker_with_cac
         image_id=image.id,
         polygon="[[0, 0], [0, 1000], [1000, 1000], [1000, 0], [0, 0]]",
     )
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/bulk/",
-        status=200,
-        json=[{"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"}],
-    )
+
+    if batch_size > 1:
+        responses.add(
+            responses.POST,
+            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/bulk/",
+            status=200,
+            json=[
+                {"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"},
+                {"id": "5468c358-b9c4-499d-8b92-d6349c58e88d"},
+            ],
+        )
+    else:
+        for elt_id in [
+            "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "5468c358-b9c4-499d-8b92-d6349c58e88d",
+        ]:
+            responses.add(
+                responses.POST,
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/bulk/",
+                status=200,
+                json=[{"id": elt_id}],
+            )
 
     created_ids = mock_elements_worker_with_cache.create_elements(
         parent=elt,
@@ -1030,30 +1057,69 @@ def test_create_elements_cached_element(responses, mock_elements_worker_with_cac
                 "name": "0",
                 "type": "something",
                 "polygon": [[1, 1], [2, 2], [2, 1], [1, 2]],
-            }
+            },
+            {
+                "name": "1",
+                "type": "something",
+                "polygon": [[4, 4], [5, 5], [5, 4], [4, 5]],
+            },
         ],
+        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/children/bulk/",
-        ),
+        )
     ]
-    assert json.loads(responses.calls[-1].request.body) == {
-        "elements": [
-            {
-                "name": "0",
-                "type": "something",
-                "polygon": [[1, 1], [2, 2], [2, 1], [1, 2]],
-            }
-        ],
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(
+            (
+                "POST",
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/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_elt = {
+        "name": "0",
+        "type": "something",
+        "polygon": [[1, 1], [2, 2], [2, 1], [1, 2]],
+    }
+    second_elt = {
+        "name": "1",
+        "type": "something",
+        "polygon": [[4, 4], [5, 5], [5, 4], [4, 5]],
+    }
+    empty_payload = {
+        "elements": [],
         "worker_run_id": "56785678-5678-5678-5678-567856785678",
     }
-    assert created_ids == [{"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"}]
+
+    bodies = []
+    first_call_idx = None
+    if batch_size > 1:
+        first_call_idx = -1
+        bodies.append({**empty_payload, "elements": [first_elt, second_elt]})
+    else:
+        first_call_idx = -2
+        bodies.append({**empty_payload, "elements": [first_elt]})
+        bodies.append({**empty_payload, "elements": [second_elt]})
+
+    assert [
+        json.loads(bulk_call.request.body)
+        for bulk_call in responses.calls[first_call_idx:]
+    ] == bodies
+
+    assert created_ids == [
+        {"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"},
+        {"id": "5468c358-b9c4-499d-8b92-d6349c58e88d"},
+    ]
 
     # Check that created elements were properly stored in SQLite cache
     assert list(CachedElement.select().order_by(CachedElement.id)) == [
@@ -1065,11 +1131,24 @@ def test_create_elements_cached_element(responses, mock_elements_worker_with_cac
             image_id="c0fec0fe-c0fe-c0fe-c0fe-c0fec0fec0fe",
             polygon=[[1, 1], [2, 2], [2, 1], [1, 2]],
             worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
+            confidence=None,
+        ),
+        CachedElement(
+            id=UUID("5468c358-b9c4-499d-8b92-d6349c58e88d"),
+            parent_id=elt.id,
+            type="something",
+            image_id="c0fec0fe-c0fe-c0fe-c0fe-c0fec0fec0fe",
+            polygon=[[4, 4], [5, 5], [5, 4], [4, 5]],
+            worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
+            confidence=None,
         ),
     ]
 
 
-def test_create_elements(responses, mock_elements_worker_with_cache, tmp_path):
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+def test_create_elements(
+    batch_size, responses, mock_elements_worker_with_cache, tmp_path
+):
     elt = Element(
         {
             "id": "12341234-1234-1234-1234-123412341234",
@@ -1083,12 +1162,28 @@ def test_create_elements(responses, mock_elements_worker_with_cache, tmp_path):
             },
         }
     )
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/bulk/",
-        status=200,
-        json=[{"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"}],
-    )
+
+    if batch_size > 1:
+        responses.add(
+            responses.POST,
+            "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/bulk/",
+            status=200,
+            json=[
+                {"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"},
+                {"id": "5468c358-b9c4-499d-8b92-d6349c58e88d"},
+            ],
+        )
+    else:
+        for elt_id in [
+            "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "5468c358-b9c4-499d-8b92-d6349c58e88d",
+        ]:
+            responses.add(
+                responses.POST,
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/bulk/",
+                status=200,
+                json=[{"id": elt_id}],
+            )
 
     created_ids = mock_elements_worker_with_cache.create_elements(
         parent=elt,
@@ -1097,30 +1192,69 @@ def test_create_elements(responses, mock_elements_worker_with_cache, tmp_path):
                 "name": "0",
                 "type": "something",
                 "polygon": [[1, 1], [2, 2], [2, 1], [1, 2]],
-            }
+            },
+            {
+                "name": "1",
+                "type": "something",
+                "polygon": [[4, 4], [5, 5], [5, 4], [4, 5]],
+            },
         ],
+        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/children/bulk/",
-        ),
+        )
     ]
-    assert json.loads(responses.calls[-1].request.body) == {
-        "elements": [
-            {
-                "name": "0",
-                "type": "something",
-                "polygon": [[1, 1], [2, 2], [2, 1], [1, 2]],
-            }
-        ],
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(
+            (
+                "POST",
+                "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/children/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_elt = {
+        "name": "0",
+        "type": "something",
+        "polygon": [[1, 1], [2, 2], [2, 1], [1, 2]],
+    }
+    second_elt = {
+        "name": "1",
+        "type": "something",
+        "polygon": [[4, 4], [5, 5], [5, 4], [4, 5]],
+    }
+    empty_payload = {
+        "elements": [],
         "worker_run_id": "56785678-5678-5678-5678-567856785678",
     }
-    assert created_ids == [{"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"}]
+
+    bodies = []
+    first_call_idx = None
+    if batch_size > 1:
+        first_call_idx = -1
+        bodies.append({**empty_payload, "elements": [first_elt, second_elt]})
+    else:
+        first_call_idx = -2
+        bodies.append({**empty_payload, "elements": [first_elt]})
+        bodies.append({**empty_payload, "elements": [second_elt]})
+
+    assert [
+        json.loads(bulk_call.request.body)
+        for bulk_call in responses.calls[first_call_idx:]
+    ] == bodies
+
+    assert created_ids == [
+        {"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"},
+        {"id": "5468c358-b9c4-499d-8b92-d6349c58e88d"},
+    ]
 
     # Check that created elements were properly stored in SQLite cache
     assert (tmp_path / "db.sqlite").is_file()
@@ -1134,7 +1268,16 @@ def test_create_elements(responses, mock_elements_worker_with_cache, tmp_path):
             polygon=[[1, 1], [2, 2], [2, 1], [1, 2]],
             worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
             confidence=None,
-        )
+        ),
+        CachedElement(
+            id=UUID("5468c358-b9c4-499d-8b92-d6349c58e88d"),
+            parent_id=UUID("12341234-1234-1234-1234-123412341234"),
+            type="something",
+            image_id="c0fec0fe-c0fe-c0fe-c0fe-c0fec0fec0fe",
+            polygon=[[4, 4], [5, 5], [5, 4], [4, 5]],
+            worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
+            confidence=None,
+        ),
     ]
 
 
@@ -1261,9 +1404,9 @@ def test_create_elements_integrity_error(
         {"id": "497f6eca-6276-4993-bfeb-53cbbbba6f08"},
     ]
 
-    assert len(caplog.records) == 1
-    assert caplog.records[0].levelname == "WARNING"
-    assert caplog.records[0].message.startswith(
+    assert len(caplog.records) == 3
+    assert caplog.records[-1].levelname == "WARNING"
+    assert caplog.records[-1].message.startswith(
         "Couldn't save created elements in local cache:"
     )
 

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