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

tests/test_elements_worker/{test_entity_create.py → test_entities.py}

@@ -3,9 +3,9 @@ import re
 from uuid import UUID
 
 import pytest
+from apistar.exceptions import ErrorResponse
 from responses import matchers
 
-from arkindex.exceptions import ErrorResponse
 from arkindex_worker.cache import (
     CachedElement,
     CachedEntity,
@@ -13,6 +13,7 @@ from arkindex_worker.cache import (
     CachedTranscriptionEntity,
 )
 from arkindex_worker.models import Transcription
+from arkindex_worker.worker.entity import MissingEntityType
 from arkindex_worker.worker.transcription import TextOrientation
 from tests import CORPUS_ID
 
@@ -684,6 +685,158 @@ def test_create_transcription_entity_with_confidence_with_cache(
     ]
 
 
+def test_list_transcription_entities_deprecation(fake_dummy_worker):
+    transcription = Transcription({"id": "fake_transcription_id"})
+    worker_version = "worker_version_id"
+    fake_dummy_worker.api_client.add_response(
+        "ListTranscriptionEntities",
+        id=transcription.id,
+        worker_version=worker_version,
+        response={"id": "entity_id"},
+    )
+    with pytest.deprecated_call(
+        match="`worker_version` usage is deprecated. Consider using `worker_run` instead."
+    ):
+        assert fake_dummy_worker.list_transcription_entities(
+            transcription, worker_version=worker_version
+        ) == {"id": "entity_id"}
+
+    assert len(fake_dummy_worker.api_client.history) == 1
+    assert len(fake_dummy_worker.api_client.responses) == 0
+
+
+def test_list_transcription_entities(fake_dummy_worker):
+    transcription = Transcription({"id": "fake_transcription_id"})
+    worker_run = "worker_run_id"
+    fake_dummy_worker.api_client.add_response(
+        "ListTranscriptionEntities",
+        id=transcription.id,
+        worker_run=worker_run,
+        response={"id": "entity_id"},
+    )
+    assert fake_dummy_worker.list_transcription_entities(
+        transcription, worker_run=worker_run
+    ) == {"id": "entity_id"}
+
+    assert len(fake_dummy_worker.api_client.history) == 1
+    assert len(fake_dummy_worker.api_client.responses) == 0
+
+
+def test_list_corpus_entities(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/entities/",
+        json={
+            "count": 1,
+            "next": None,
+            "results": [
+                {
+                    "id": "fake_entity_id",
+                }
+            ],
+        },
+    )
+
+    mock_elements_worker.list_corpus_entities()
+
+    assert mock_elements_worker.entities == {
+        "fake_entity_id": {
+            "id": "fake_entity_id",
+        }
+    }
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        (
+            "GET",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/entities/",
+        ),
+    ]
+
+
+@pytest.mark.parametrize("wrong_name", [1234, 12.5])
+def test_list_corpus_entities_wrong_name(mock_elements_worker, wrong_name):
+    with pytest.raises(AssertionError, match="name should be of type str"):
+        mock_elements_worker.list_corpus_entities(name=wrong_name)
+
+
+@pytest.mark.parametrize("wrong_parent", [{"id": "element_id"}, 12.5, "blabla"])
+def test_list_corpus_entities_wrong_parent(mock_elements_worker, wrong_parent):
+    with pytest.raises(AssertionError, match="parent should be of type Element"):
+        mock_elements_worker.list_corpus_entities(parent=wrong_parent)
+
+
+def test_check_required_entity_types(responses, mock_elements_worker):
+    # Set one entity type
+    mock_elements_worker.entity_types = {"person": "person-entity-type-id"}
+
+    checked_types = ["person", "new-entity"]
+
+    # Call to create new entity type
+    responses.add(
+        responses.POST,
+        "http://testserver/api/v1/entity/types/",
+        status=200,
+        match=[
+            matchers.json_params_matcher(
+                {
+                    "name": "new-entity",
+                    "corpus": CORPUS_ID,
+                }
+            )
+        ],
+        json={
+            "id": "new-entity-id",
+            "corpus": CORPUS_ID,
+            "name": "new-entity",
+            "color": "ffd1b3",
+        },
+    )
+
+    mock_elements_worker.check_required_entity_types(
+        entity_types=checked_types,
+    )
+
+    # Make sure the entity_types entry has been updated
+    assert mock_elements_worker.entity_types == {
+        "person": "person-entity-type-id",
+        "new-entity": "new-entity-id",
+    }
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        (
+            "POST",
+            "http://testserver/api/v1/entity/types/",
+        ),
+    ]
+
+
+def test_check_required_entity_types_no_creation_allowed(
+    responses, mock_elements_worker
+):
+    # Set one entity type
+    mock_elements_worker.entity_types = {"person": "person-entity-type-id"}
+
+    checked_types = ["person", "new-entity"]
+
+    with pytest.raises(
+        MissingEntityType, match="Entity type `new-entity` was not in the corpus."
+    ):
+        mock_elements_worker.check_required_entity_types(
+            entity_types=checked_types, create_missing=False
+        )
+
+    assert len(responses.calls) == len(BASE_API_CALLS)
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS
+
+
 @pytest.mark.parametrize("transcription", [None, "not a transcription", 1])
 def test_create_transcription_entities_wrong_transcription(
     mock_elements_worker, transcription
@@ -837,7 +990,6 @@ def test_create_transcription_entities_wrong_entity(
 
 def test_create_transcription_entities(responses, mock_elements_worker):
     transcription = Transcription(id="transcription-id")
-
     # Call to Transcription entities creation in bulk
     responses.add(
         responses.POST,
@@ -854,14 +1006,7 @@ 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,
-                        },
+                        }
                     ],
                 }
             )
@@ -870,12 +1015,8 @@ def test_create_transcription_entities(responses, mock_elements_worker):
             "entities": [
                 {
                     "transcription_entity_id": "transc-entity-id",
-                    "entity_id": "entity-id1",
-                },
-                {
-                    "transcription_entity_id": "transc-entity-id",
-                    "entity_id": "entity-id2",
-                },
+                    "entity_id": "entity-id",
+                }
             ]
         },
     )
@@ -893,18 +1034,11 @@ 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,
-            },
+            }
         ],
     )
 
-    assert len(created_objects) == 2
+    assert len(created_objects) == 1
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
@@ -913,5 +1047,5 @@ def test_create_transcription_entities(responses, mock_elements_worker):
         (
             "POST",
             "http://testserver/api/v1/transcription/transcription-id/entities/bulk/",
-        )
+        ),
     ]

tests/test_elements_worker/test_image.py

@@ -1,8 +1,7 @@
 import json
 
 import pytest
-
-from arkindex.exceptions import ErrorResponse
+from apistar.exceptions import ErrorResponse
 
 from . import BASE_API_CALLS
 

tests/test_elements_worker/test_metadata.py

@@ -2,12 +2,11 @@ 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
@@ -229,267 +228,143 @@ def test_create_metadata_cached_element(responses, mock_elements_worker_with_cac
     assert metadata_id == "12345678-1234-1234-1234-123456789123"
 
 
-@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
-def test_create_metadata_bulk(batch_size, responses, mock_elements_worker):
+@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):
     element = Element({"id": "12341234-1234-1234-1234-123412341234"})
-
-    metadata_list = [
-        {"type": MetaType.Text, "name": "fake_name", "value": "fake_value"},
-        {
-            "type": MetaType.Text,
-            "name": "Year",
-            "value": "2024",
-            "entity_id": "entity_id",
+    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"),
+                }
+            ],
         },
-    ]
-    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, batch_size=batch_size
+        element, metadata_list
     )
 
-    bulk_api_calls = [
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
         (
             "POST",
             "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        )
+        ),
     ]
-    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 == [
+    assert json.loads(responses.calls[-1].request.body)["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,
-        },
+            "entity_id": metadata_list[0].get("entity_id"),
+        }
+    ]
+    assert created_metadata_list == [
         {
-            "id": "fake_metadata_id2",
-            "type": metadata_list[1]["type"].value,
-            "name": metadata_list[1]["name"],
-            "value": metadata_list[1]["value"],
+            "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[1]["entity_id"],
-        },
+            "entity_id": metadata_list[0].get("entity_id"),
+        }
     ]
 
 
-@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+@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_cached_element(
-    batch_size, responses, mock_elements_worker_with_cache
+    responses, mock_elements_worker_with_cache, metadata_list
 ):
     element = CachedElement.create(
         id="12341234-1234-1234-1234-123412341234", type="thing"
     )
-
-    metadata_list = [
-        {"type": MetaType.Text, "name": "fake_name", "value": "fake_value"},
-        {
-            "type": MetaType.Text,
-            "name": "Year",
-            "value": "2024",
-            "entity_id": "entity_id",
+    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"),
+                }
+            ],
         },
-    ]
-    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, batch_size=batch_size
+        element, metadata_list
     )
 
-    bulk_api_calls = [
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
         (
             "POST",
             "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/metadata/bulk/",
-        )
+        ),
     ]
-    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 == [
+    assert json.loads(responses.calls[-1].request.body)["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,
-        },
+            "entity_id": metadata_list[0].get("entity_id"),
+        }
+    ]
+    assert created_metadata_list == [
         {
-            "id": "fake_metadata_id2",
-            "type": metadata_list[1]["type"].value,
-            "name": metadata_list[1]["name"],
-            "value": metadata_list[1]["value"],
+            "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[1]["entity_id"],
-        },
+            "entity_id": metadata_list[0].get("entity_id"),
+        }
     ]
 
 

tests/test_elements_worker/test_task.py

@@ -1,8 +1,8 @@
 import uuid
 
 import pytest
+from apistar.exceptions import ErrorResponse
 
-from arkindex.exceptions import ErrorResponse
 from arkindex_worker.models import Artifact
 from tests import FIXTURES_DIR
 from tests.test_elements_worker import BASE_API_CALLS

tests/test_elements_worker/test_training.py

@@ -8,7 +8,7 @@ from arkindex_worker.worker import BaseWorker
 from arkindex_worker.worker.training import TrainingMixin, create_archive
 
 
-@pytest.fixture
+@pytest.fixture()
 def mock_training_worker(monkeypatch):
     class TrainingWorker(BaseWorker, TrainingMixin):
         """
@@ -22,7 +22,7 @@ def mock_training_worker(monkeypatch):
     return training_worker
 
 
-@pytest.fixture
+@pytest.fixture()
 def default_model_version():
     return {
         "id": "model_version_id",