Users can now update assets with attached media that contains images. (#649)

* do not restrict connectors to not uploading binaries

* Allow users to update assets if they do not change binaries

* Put back accidentally removed unit tests

Author: PGijsbers

src/database/model/ai_asset/distribution.py

@@ -3,7 +3,7 @@
 from typing import Type
 
 from fastapi.encoders import jsonable_encoder
-from pydantic import create_model
+from pydantic import create_model, validator
 from sqlalchemy import Column, ForeignKey, String, LargeBinary
 from sqlmodel import Field
 
@@ -64,6 +64,14 @@ class DistributionBase(AIoDConceptBase):
         sa_column=Column(LargeBinary),
     )
 
+    @validator("binary_blob", pre=True, always=True)
+    def decode_string_to_bytes(cls, v) -> bytes | None:
+        if v is None or isinstance(v, bytes):
+            return v
+        if not isinstance(v, str):
+            raise TypeError("`binary_blob` can only be bytes, str, or None.")
+        return base64.b64decode(v)
+
     def dict(self, *args, **kwargs):
         # Defining it as a `Config` does not work for some reason.
         item_dict = super().dict(*args, **kwargs)

src/routers/resource_router.py

@@ -544,7 +544,17 @@ def put_resource(
             with DbSession() as session:
                 try:
                     resource: Any = self._retrieve_resource(session, identifier)
-                    _raise_if_contains_binary_blob(resource_create_instance)
+                    if not user.is_connector:
+                        if hasattr(resource_create_instance, "media"):
+                            if not resource_create_instance.media:  # type: ignore[attr-defined]
+                                # This does create the problem that a user cannot remove all media through this endpoint :/
+                                resource_create_instance.media = resource.media  # type: ignore[attr-defined]
+                            elif set(m.binary_blob for m in resource_create_instance.media) != set(  # type: ignore[attr-defined]
+                                m.binary_blob for m in resource.media
+                            ):
+                                _raise_if_contains_binary_blob(resource_create_instance)
+                        else:
+                            _raise_if_contains_binary_blob(resource_create_instance)
                     if not (
                         user_can_write(user, resource.aiod_entry)
                         or user.has_role(f"update_{self.resource_name_plural}")

src/tests/routers/resource_routers/test_router_organisation.py

@@ -1,6 +1,7 @@
 import copy
 from unittest.mock import Mock
 
+from fastapi.encoders import jsonable_encoder
 from starlette.testclient import TestClient
 
 from database.model.agent.contact import Contact
@@ -351,3 +352,82 @@ def test_organisation_delete_image(
         )
         assert second_delete_response.status_code == HTTPStatus.NOT_FOUND
         assert "No image with the name" in second_delete_response.json()["detail"]
+
+
+def test_organisation_put_without_media_keeps_media(
+        client: TestClient,
+        organisation: Organisation,
+):
+    organisation.media = []
+    identifier = register_asset(organisation)
+
+    fake_image = io.BytesIO(b"\x89PNG\r\n\x1a\n...")  # fake PNG bytes
+    fake_image.name = "logo.png"
+
+    with logged_in_user():
+        response = client.post(
+            f"/organisations/{identifier}/image",
+            params={"name": "logo"},
+            files={"file": ("logo.png", fake_image, "image/png")},
+            headers={"Authorization": "Fake token"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+
+        organisation.name = "new name"
+        response = client.put(
+            f"/organisations/{identifier}",
+            json=jsonable_encoder(organisation.dict()),
+            headers={"Authorization": "Fake token"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+        response = client.get(
+            f"/organisations/{identifier}",
+            headers={"Authorization": "Fake token"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+        assert response.json()["name"] == "new name", response.json()
+        assert response.json()["media"], response.json()
+
+
+def test_organisation_put_with_media_keeps_media_if_no_new_binary(
+        client: TestClient,
+        organisation: Organisation,
+):
+    organisation.media = []
+    identifier = register_asset(organisation)
+
+    fake_image = io.BytesIO(b"\x89PNG\r\n\x1a\n...")  # fake PNG bytes
+    fake_image.name = "logo.png"
+
+    with logged_in_user():
+        response = client.post(
+            f"/organisations/{identifier}/image",
+            params={"name": "logo"},
+            files={"file": ("logo.png", fake_image, "image/png")},
+            headers={"Authorization": "Fake token"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+
+        response = client.get(
+            f"/organisations/{identifier}?get_image=true",
+            headers={"Authorization": "Fake token"},
+        )
+        org = response.json()
+        del org["aiod_entry"]
+        org["media"].append(
+            {"name": "foo", "binary_blob": "bar="},
+        )
+        response = client.put(
+            f"/organisations/{identifier}",
+            json=org,
+            headers={"Authorization": "Fake token"},
+        )
+        assert response.status_code == HTTPStatus.BAD_REQUEST, "No new binary may be added through a PUT request"
+
+        org["media"].pop()
+        response = client.put(
+            f"/organisations/{identifier}",
+            json=org,
+            headers={"Authorization": "Fake token"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()