Check before updating image, better tests

arkindex_worker/worker/element.py

@@ -343,7 +343,8 @@ class ElementMixin(object):
                 payload["polygon"] = str(payload["polygon"])
 
             # Retrieve the right image
-            payload["image"] = CachedImage.get_by_id(payload["image"])
+            if "image" in payload:
+                payload["image"] = CachedImage.get_by_id(payload["image"])
 
             CachedElement.update(**payload).where(
                 CachedElement.id == element.id

tests/test_elements_worker/test_elements.py

@@ -1306,8 +1306,37 @@ def test_update_element_api_error(responses, mock_elements_worker):
     ]
 
 
+@pytest.mark.parametrize(
+    "payload",
+    (
+        (
+            {
+                "polygon": [[10, 10], [20, 20], [20, 10], [10, 20]],
+                "confidence": None,
+            }
+        ),
+        (
+            {
+                "rotation_angle": 45,
+                "mirrored": False,
+            }
+        ),
+        (
+            {
+                "polygon": [[10, 10], [20, 20], [20, 10], [10, 20]],
+                "confidence": None,
+                "rotation_angle": 45,
+                "mirrored": False,
+            }
+        ),
+    ),
+)
 def test_update_element(
-    responses, mock_elements_worker_with_cache, mock_cached_elements, mock_cached_images
+    responses,
+    mock_elements_worker_with_cache,
+    mock_cached_elements,
+    mock_cached_images,
+    payload,
 ):
     elt = CachedElement.select().first()
     new_image = CachedImage.select().first()
@@ -1315,11 +1344,8 @@ def test_update_element(
     elt_response = {
         "type": "new type",
         "name": "new name",
-        "polygon": [[10, 10], [20, 20], [20, 10], [10, 20]],
-        "confidence": None,
-        "rotation_angle": 45,
-        "mirrored": False,
         "image": str(new_image.id),
+        **payload,
     }
     responses.add(
         responses.PUT,
@@ -1347,12 +1373,16 @@ def test_update_element(
     assert element_update_response == elt_response
 
     cached_element = CachedElement.get(CachedElement.id == elt.id)
+    # Always present in payload
     assert cached_element.type == elt_response["type"]
-    assert cached_element.polygon == str(elt_response["polygon"])
-    assert cached_element.confidence == elt_response["confidence"]
-    assert cached_element.rotation_angle == elt_response["rotation_angle"]
-    assert cached_element.mirrored == elt_response["mirrored"]
     assert str(cached_element.image_id) == elt_response["image"]
+    # Optional params
+    if "polygon" in payload:
+        # Cast to string as this is the only difference compared to model
+        elt_response["polygon"] = str(elt_response["polygon"])
+
+    for param in payload:
+        assert getattr(cached_element, param) == elt_response[param]
 
 
 def test_update_element_confidence(