Add error handling for an invalid image file and add image processing…

… tests #29

object_storage_api/core/exceptions.py

@@ -44,3 +44,12 @@ class InvalidObjectIdError(DatabaseError):
 
     status_code = 422
     response_detail = "Invalid ID given"
+
+
+class InvalidImageFileError(BaseAPIException):
+    """
+    The provided image file is not valid.
+    """
+
+    status_code = 422
+    response_detail = "File given is not a valid image"

object_storage_api/core/image.py

@@ -7,9 +7,10 @@
 from io import BytesIO
 
 from fastapi import UploadFile
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 
 from object_storage_api.core.config import config
+from object_storage_api.core.exceptions import InvalidImageFileError
 
 logger = logging.getLogger()
 
@@ -22,11 +23,18 @@ def generate_thumbnail_base64_str(uploaded_image_file: UploadFile) -> str:
 
     :param uploaded_image_file: Uploaded image file.
     :return: Base64 encoded string of the thumbnail
+    :raises: InvalidImageFileError if the given image file cannot be processed due to being invalid in some way.
     """
 
     logger.debug("Generating thumbnail for uploaded image file")
 
-    pillow_image = Image.open(uploaded_image_file.file)
+    # Image may fail to open if the file is either not an image or is invalid in some other way
+    try:
+        pillow_image = Image.open(uploaded_image_file.file)
+    except UnidentifiedImageError as exc:
+        raise InvalidImageFileError(
+            f"The uploaded file '{uploaded_image_file.filename}' could not be opened by Pillow"
+        ) from exc
 
     pillow_image.thumbnail(
         (image_config.thumbnail_max_size_pixels, image_config.thumbnail_max_size_pixels),

test/e2e/test_image.py

@@ -29,17 +29,18 @@ def setup(self, test_client):
 
         self.test_client = test_client
 
-    def post_image(self, image_post_metadata_data: dict) -> Optional[str]:
+    def post_image(self, image_post_metadata_data: dict, file_name: str) -> Optional[str]:
         """
         Posts an image with the given metadata and a test image file and returns the id of the created image if
         successful.
 
         :param image_post_metadata_data: Dictionary containing the image metadata data as would be required for an
                                          `ImagePostMetadataSchema`.
+        :param file_name: File name of the image to upload (relative to the 'test/files' directory).
         :return: ID of the created image (or `None` if not successful).
         """
 
-        with open("test/e2e/files/image.jpg", mode="rb") as file:
+        with open(f"test/files/{file_name}", mode="rb") as file:
             self._post_response_image = self.test_client.post(
                 "/images", data={**image_post_metadata_data}, files={"upload_file": file}
             )
@@ -74,17 +75,23 @@ class TestCreate(CreateDSL):
     def test_create_with_only_required_values_provided(self):
         """Test creating an image with only required values provided."""
 
-        self.post_image(IMAGE_POST_METADATA_DATA_REQUIRED_VALUES_ONLY)
+        self.post_image(IMAGE_POST_METADATA_DATA_REQUIRED_VALUES_ONLY, "image.jpg")
         self.check_post_image_success(IMAGE_GET_DATA_REQUIRED_VALUES_ONLY)
 
     def test_create_with_all_values_provided(self):
         """Test creating an image with all values provided."""
 
-        self.post_image(IMAGE_POST_METADATA_DATA_ALL_VALUES)
+        self.post_image(IMAGE_POST_METADATA_DATA_ALL_VALUES, "image.jpg")
         self.check_post_image_success(IMAGE_GET_DATA_ALL_VALUES)
 
     def test_create_with_invalid_entity_id(self):
         """Test creating an image with an invalid `entity_id`."""
 
-        self.post_image({**IMAGE_POST_METADATA_DATA_REQUIRED_VALUES_ONLY, "entity_id": "invalid-id"})
+        self.post_image({**IMAGE_POST_METADATA_DATA_REQUIRED_VALUES_ONLY, "entity_id": "invalid-id"}, "image.jpg")
         self.check_post_image_failed_with_detail(422, "Invalid `entity_id` given")
+
+    def test_create_with_invalid_image_file(self):
+        """Test creating an image with an invalid image file."""
+
+        self.post_image(IMAGE_POST_METADATA_DATA_REQUIRED_VALUES_ONLY, "invalid_image.jpg")
+        self.check_post_image_failed_with_detail(422, "File given is not a valid image")

test/unit/core/test_image.py

@@ -0,0 +1,32 @@
+"""
+Unit tests for image processing functions.
+"""
+
+import pytest
+from fastapi import UploadFile
+
+from object_storage_api.core.exceptions import InvalidImageFileError
+from object_storage_api.core.image import generate_thumbnail_base64_str
+
+
+class TestGenerateThumbnailBase64Str:
+    """Tests for the `generate_thumbnail_base64_str` method."""
+
+    def test_with_valid_image(self):
+        """Tests `generate_thumbnail_base64_str` with a valid image file provided."""
+
+        with open("test/files/image.jpg", "rb") as file:
+            uploaded_image_file = UploadFile(file, filename="image.jpg")
+            result = generate_thumbnail_base64_str(uploaded_image_file)
+
+        assert result == "UklGRjQAAABXRUJQVlA4ICgAAADQAQCdASoCAAEAAUAmJYwCdAEO/gOOAAD+qlQWHDxhNJOjVlqIb8AA"
+
+    def test_with_invalid_image(self):
+        """Tests `generate_thumbnail_base64_str` with an invalid image file provided."""
+
+        with open("test/files/invalid_image.jpg", "rb") as file:
+            uploaded_image_file = UploadFile(file, filename="image.jpg")
+            with pytest.raises(InvalidImageFileError) as exc:
+                generate_thumbnail_base64_str(uploaded_image_file)
+
+            assert str(exc.value) == f"The uploaded file '{uploaded_image_file.filename}' could not be opened by Pillow"