Add list_element_children function

arkindex_worker/worker.py

@@ -716,3 +716,69 @@ class ElementsWorker(BaseWorker):
         )
 
         return transcriptions
+
+    def list_element_children(
+        self,
+        element,
+        best_class=None,
+        folder=None,
+        name=None,
+        recursive=None,
+        type=None,
+        with_best_classes=None,
+        with_corpus=None,
+        with_has_children=None,
+        with_zone=None,
+        worker_version=None,
+    ):
+        """
+        List children of an element
+        """
+        assert element and isinstance(
+            element, Element
+        ), "element shouldn't be null and should be of type Element"
+        query_params = {}
+        if best_class is not None:
+            assert isinstance(best_class, str) or isinstance(
+                best_class, bool
+            ), "best_class should be of type str or bool"
+            query_params["best_class"] = best_class
+        if folder is not None:
+            assert isinstance(folder, bool), "folder should be of type bool"
+            query_params["folder"] = folder
+        if name:
+            assert isinstance(name, str), "name should be of type str"
+            query_params["name"] = name
+        if recursive is not None:
+            assert isinstance(recursive, bool), "recursive should be of type bool"
+            query_params["recursive"] = recursive
+        if type:
+            assert isinstance(type, str), "type should be of type str"
+            query_params["type"] = type
+        if with_best_classes is not None:
+            assert isinstance(
+                with_best_classes, bool
+            ), "with_best_classes should be of type bool"
+            query_params["with_best_classes"] = with_best_classes
+        if with_corpus is not None:
+            assert isinstance(with_corpus, bool), "with_corpus should be of type bool"
+            query_params["with_corpus"] = with_corpus
+        if with_has_children is not None:
+            assert isinstance(
+                with_has_children, bool
+            ), "with_has_children should be of type bool"
+            query_params["with_has_children"] = with_has_children
+        if with_zone is not None:
+            assert isinstance(with_zone, bool), "with_zone should be of type bool"
+            query_params["with_zone"] = with_zone
+        if worker_version:
+            assert isinstance(
+                worker_version, str
+            ), "worker_version should be of type str"
+            query_params["worker_version"] = worker_version
+
+        children = self.api_client.paginate(
+            "ListElementChildren", id=element.id, **query_params
+        )
+
+        return children

tests/test_elements_worker.py

@@ -2163,3 +2163,208 @@ def test_list_transcriptions(responses, mock_elements_worker):
         "http://testserver/api/v1/workers/versions/12341234-1234-1234-1234-123412341234/",
         "http://testserver/api/v1/element/12341234-1234-1234-1234-123412341234/transcriptions/",
     ]
+
+
+def test_list_element_children_wrong_element(mock_elements_worker):
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(element=None)
+    assert str(e.value) == "element shouldn't be null and should be of type Element"
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(element="not element type")
+    assert str(e.value) == "element shouldn't be null and should be of type Element"
+
+
+def test_list_element_children_wrong_best_class(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            best_class=1234,
+        )
+    assert str(e.value) == "best_class should be of type str or bool"
+
+
+def test_list_element_children_wrong_folder(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            folder="not bool",
+        )
+    assert str(e.value) == "folder should be of type bool"
+
+
+def test_list_element_children_wrong_name(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            name=1234,
+        )
+    assert str(e.value) == "name should be of type str"
+
+
+def test_list_element_children_wrong_recursive(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            recursive="not bool",
+        )
+    assert str(e.value) == "recursive should be of type bool"
+
+
+def test_list_element_children_wrong_type(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            type=1234,
+        )
+    assert str(e.value) == "type should be of type str"
+
+
+def test_list_element_children_wrong_with_best_classes(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            with_best_classes="not bool",
+        )
+    assert str(e.value) == "with_best_classes should be of type bool"
+
+
+def test_list_element_children_wrong_with_corpus(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            with_corpus="not bool",
+        )
+    assert str(e.value) == "with_corpus should be of type bool"
+
+
+def test_list_element_children_wrong_with_has_children(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            with_has_children="not bool",
+        )
+    assert str(e.value) == "with_has_children should be of type bool"
+
+
+def test_list_element_children_wrong_with_zone(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            with_zone="not bool",
+        )
+    assert str(e.value) == "with_zone should be of type bool"
+
+
+def test_list_element_children_wrong_worker_version(mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.list_element_children(
+            element=elt,
+            worker_version=1234,
+        )
+    assert str(e.value) == "worker_version should be of type str"
+
+
+def test_list_element_children_api_error(responses, mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+        status=500,
+    )
+
+    with pytest.raises(
+        Exception, match="Stopping pagination as data will be incomplete"
+    ):
+        next(mock_elements_worker.list_element_children(element=elt))
+
+    assert len(responses.calls) == 6
+    assert [call.request.url for call in responses.calls] == [
+        "http://testserver/api/v1/workers/versions/12341234-1234-1234-1234-123412341234/",
+        # We do 5 retries
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+    ]
+
+
+def test_list_element_children(responses, mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+    expected_children = [
+        {
+            "id": "0000",
+            "type": "page",
+            "name": "Test",
+            "corpus": {},
+            "thumbnail_url": None,
+            "zone": {},
+            "best_classes": None,
+            "has_children": None,
+            "worker_version_id": None,
+        },
+        {
+            "id": "1111",
+            "type": "page",
+            "name": "Test 2",
+            "corpus": {},
+            "thumbnail_url": None,
+            "zone": {},
+            "best_classes": None,
+            "has_children": None,
+            "worker_version_id": None,
+        },
+        {
+            "id": "2222",
+            "type": "page",
+            "name": "Test 3",
+            "corpus": {},
+            "thumbnail_url": None,
+            "zone": {},
+            "best_classes": None,
+            "has_children": None,
+            "worker_version_id": None,
+        },
+    ]
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+        status=200,
+        json={
+            "count": 3,
+            "next": None,
+            "results": expected_children,
+        },
+    )
+
+    for idx, child in enumerate(
+        mock_elements_worker.list_element_children(element=elt)
+    ):
+        assert child == expected_children[idx]
+
+    assert len(responses.calls) == 2
+    assert [call.request.url for call in responses.calls] == [
+        "http://testserver/api/v1/workers/versions/12341234-1234-1234-1234-123412341234/",
+        "http://testserver/api/v1/elements/12341234-1234-1234-1234-123412341234/children/",
+    ]