From d55250aac17d6f796bc36e35d15df3d947bcacf2 Mon Sep 17 00:00:00 2001 From: Joel Davies Date: Tue, 8 Oct 2024 09:37:30 +0000 Subject: [PATCH] Add error handling for an invalid image file and add image processing tests #29 --- object_storage_api/core/exceptions.py | 9 ++++++++ object_storage_api/core/image.py | 12 ++++++++-- test/e2e/test_image.py | 17 ++++++++++---- test/{e2e => }/files/image.jpg | Bin test/files/invalid_image.jpg | 1 + test/unit/core/test_image.py | 32 ++++++++++++++++++++++++++ 6 files changed, 64 insertions(+), 7 deletions(-) rename test/{e2e => }/files/image.jpg (100%) create mode 100644 test/files/invalid_image.jpg create mode 100644 test/unit/core/test_image.py diff --git a/object_storage_api/core/exceptions.py b/object_storage_api/core/exceptions.py index d2e562a..095c051 100644 --- a/object_storage_api/core/exceptions.py +++ b/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" diff --git a/object_storage_api/core/image.py b/object_storage_api/core/image.py index 040f58e..b8fa758 100644 --- a/object_storage_api/core/image.py +++ b/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), diff --git a/test/e2e/test_image.py b/test/e2e/test_image.py index c5746a0..5596fbd 100644 --- a/test/e2e/test_image.py +++ b/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") diff --git a/test/e2e/files/image.jpg b/test/files/image.jpg similarity index 100% rename from test/e2e/files/image.jpg rename to test/files/image.jpg diff --git a/test/files/invalid_image.jpg b/test/files/invalid_image.jpg new file mode 100644 index 0000000..8bf5d97 --- /dev/null +++ b/test/files/invalid_image.jpg @@ -0,0 +1 @@ +Not image data \ No newline at end of file diff --git a/test/unit/core/test_image.py b/test/unit/core/test_image.py new file mode 100644 index 0000000..ebdb981 --- /dev/null +++ b/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"