Add an API helper to update an element

arkindex_worker/worker/element.py

@@ -275,6 +275,70 @@ class ElementMixin(object):
 
         return created_ids
 
+    def update_element(
+        self,
+        element: Union[Element, CachedElement],
+        type: Optional[str] = None,
+        name: Optional[str] = None,
+        polygon: Optional[List[List[Union[int, float]]]] = None,
+        confidence: Optional[float] = None,
+    ) -> dict:
+        """
+        Partially update an element through the API.
+
+        :param element: The element to update.
+        :param type: Optional new slug type of the element.
+        :param name: Optional new name of the element.
+        :param polygon: Optional new polygon of the element.
+        :param confidence: Optional new confidence score, between 0.0 and 1.0.
+        :returns: A dict from the ``PartialUpdateElement`` API endpoint,
+        """
+        assert element and isinstance(
+            element, (Element, CachedElement)
+        ), "element shouldn't be null and should be an Element or CachedElement"
+        assert type is None or isinstance(type, str), "type should be None or a str"
+        assert name is None or isinstance(name, str), "name should be None or a str"
+        assert polygon is None or isinstance(
+            polygon, list
+        ), "polygon should be None or a list"
+        if polygon:
+            assert len(polygon) >= 3, "polygon should have at least three points"
+            assert all(
+                isinstance(point, list) and len(point) == 2 for point in polygon
+            ), "polygon points should be lists of two items"
+            assert all(
+                isinstance(coord, (int, float)) for point in polygon for coord in point
+            ), "polygon points should be lists of two numbers"
+        assert confidence is None or (
+            isinstance(confidence, float) and 0 <= confidence <= 1
+        ), "confidence should be None or a float in [0..1] range"
+
+        if self.is_read_only:
+            logger.warning("Cannot update element as this worker is in read-only mode")
+            return
+
+        updated_element = self.request(
+            "PartialUpdateElement",
+            id=element.id,
+            body={
+                "type": type,
+                "name": name,
+                "polygon": polygon,
+                "confidence": confidence,
+            },
+        )
+
+        if self.use_cache:
+            CachedElement.update(
+                {
+                    CachedElement.type: type,
+                    CachedElement.polygon: str(polygon),
+                    CachedElement.confidence: confidence,
+                }
+            ).where(CachedElement.id == element.id).execute()
+
+        return updated_element
+
     def list_element_children(
         self,
         element: Union[Element, CachedElement],

tests/test_elements_worker/test_elements.py

@@ -1210,6 +1210,202 @@ def test_create_elements_integrity_error(
     assert list(CachedElement.select()) == []
 
 
+def test_update_element_wrong_element(mock_elements_worker):
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=None,
+        )
+    assert (
+        str(e.value)
+        == "element shouldn't be null and should be an Element or CachedElement"
+    )
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element="not element type",
+        )
+    assert (
+        str(e.value)
+        == "element shouldn't be null and should be an Element or CachedElement"
+    )
+
+
+def test_update_element_wrong_type(mock_elements_worker):
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=Element({"zone": None}),
+            type=1234,
+        )
+    assert str(e.value) == "type should be None or a str"
+
+
+def test_update_element_wrong_name(mock_elements_worker):
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=Element({"zone": None}),
+            name=1234,
+        )
+    assert str(e.value) == "name should be None or a str"
+
+
+def test_update_element_wrong_polygon(mock_elements_worker):
+    elt = Element({"zone": None})
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=elt,
+            polygon="not a polygon",
+        )
+    assert str(e.value) == "polygon should be None or a list"
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=elt,
+            polygon=[[1, 1], [2, 2]],
+        )
+    assert str(e.value) == "polygon should have at least three points"
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=elt,
+            polygon=[[1, 1, 1], [2, 2, 1], [2, 1, 1], [1, 2, 1]],
+        )
+    assert str(e.value) == "polygon points should be lists of two items"
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=elt,
+            polygon=[[1], [2], [2], [1]],
+        )
+    assert str(e.value) == "polygon points should be lists of two items"
+
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=elt,
+            polygon=[["not a coord", 1], [2, 2], [2, 1], [1, 2]],
+        )
+    assert str(e.value) == "polygon points should be lists of two numbers"
+
+
+@pytest.mark.parametrize("confidence", ["lol", "0.2", -1.0, 1.42, float("inf")])
+def test_update_element_wrong_confidence(mock_elements_worker, confidence):
+    with pytest.raises(AssertionError) as e:
+        mock_elements_worker.update_element(
+            element=Element({"zone": None}),
+            confidence=confidence,
+        )
+    assert str(e.value) == "confidence should be None or a float in [0..1] range"
+
+
+def test_update_element_api_error(responses, mock_elements_worker):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+    responses.add(
+        responses.PATCH,
+        f"http://testserver/api/v1/element/{elt.id}/",
+        status=500,
+    )
+
+    with pytest.raises(ErrorResponse):
+        mock_elements_worker.update_element(
+            element=elt,
+            type="something",
+            name="0",
+            polygon=[[1, 1], [2, 2], [2, 1], [1, 2]],
+        )
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        # We retry 5 times the API call
+        ("PATCH", f"http://testserver/api/v1/element/{elt.id}/"),
+        ("PATCH", f"http://testserver/api/v1/element/{elt.id}/"),
+        ("PATCH", f"http://testserver/api/v1/element/{elt.id}/"),
+        ("PATCH", f"http://testserver/api/v1/element/{elt.id}/"),
+        ("PATCH", f"http://testserver/api/v1/element/{elt.id}/"),
+    ]
+
+
+def test_update_element(
+    responses, mock_elements_worker_with_cache, mock_cached_elements
+):
+    elt = CachedElement.select().first()
+    elt_response = {
+        "type": "new type",
+        "name": "new name",
+        "polygon": [[10, 10], [20, 20], [20, 10], [10, 20]],
+        "confidence": None,
+    }
+    responses.add(
+        responses.PATCH,
+        f"http://testserver/api/v1/element/{elt.id}/",
+        status=200,
+        json=elt_response,
+    )
+
+    element_update_response = mock_elements_worker_with_cache.update_element(
+        element=elt,
+        **elt_response,
+    )
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        (
+            "PATCH",
+            f"http://testserver/api/v1/element/{elt.id}/",
+        ),
+    ]
+    assert json.loads(responses.calls[-1].request.body) == elt_response
+    assert element_update_response == elt_response
+
+    cached_element = CachedElement.get(CachedElement.id == elt.id)
+    assert cached_element.type == elt_response["type"]
+    assert cached_element.polygon == str(elt_response["polygon"])
+    assert cached_element.confidence == elt_response["confidence"]
+
+
+def test_update_element_confidence(
+    responses, mock_elements_worker_with_cache, mock_cached_elements
+):
+    elt = CachedElement.select().first()
+    elt_response = {
+        "type": "new type",
+        "name": "new name",
+        "polygon": [[10, 10], [20, 20], [20, 10], [10, 20]],
+        "confidence": 0.42,
+    }
+    responses.add(
+        responses.PATCH,
+        f"http://testserver/api/v1/element/{elt.id}/",
+        status=200,
+        json=elt_response,
+    )
+
+    element_update_response = mock_elements_worker_with_cache.update_element(
+        element=elt,
+        **elt_response,
+    )
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        (
+            "PATCH",
+            f"http://testserver/api/v1/element/{elt.id}/",
+        ),
+    ]
+    assert json.loads(responses.calls[-1].request.body) == elt_response
+    assert element_update_response == elt_response
+
+    cached_element = CachedElement.get(CachedElement.id == elt.id)
+    assert cached_element.type == elt_response["type"]
+    assert cached_element.polygon == str(elt_response["polygon"])
+    assert cached_element.confidence == elt_response["confidence"]
+
+
 def test_list_element_children_wrong_element(mock_elements_worker):
     with pytest.raises(AssertionError) as e:
         mock_elements_worker.list_element_children(element=None)