diff --git a/object_storage_api/services/image.py b/object_storage_api/services/image.py index f9906d9..401fbc1 100644 --- a/object_storage_api/services/image.py +++ b/object_storage_api/services/image.py @@ -4,6 +4,7 @@ """ import logging +import mimetypes from typing import Annotated, Optional from bson import ObjectId @@ -58,6 +59,12 @@ def create(self, image_metadata: ImagePostMetadataSchema, upload_file: UploadFil # Upload the full size image to object storage object_key = self._image_store.upload(image_id, image_metadata, upload_file) + expected_file_type = mimetypes.guess_type(upload_file.filename)[0] + if expected_file_type != upload_file.content_type: + raise InvalidObjectIdError( + f"File extension `{upload_file.filename}` does not match content type `{upload_file.content_type}`" + ) + try: image_in = ImageIn( **image_metadata.model_dump(), diff --git a/test/unit/services/test_image.py b/test/unit/services/test_image.py index d39649a..d48c858 100644 --- a/test/unit/services/test_image.py +++ b/test/unit/services/test_image.py @@ -62,16 +62,18 @@ class CreateDSL(ImageServiceDSL): _created_image: ImageMetadataSchema _create_exception: pytest.ExceptionInfo - def mock_create(self, image_post_metadata_data: dict) -> None: + def mock_create(self, image_post_metadata_data: dict, filename: str) -> None: """ Mocks repo & store methods appropriately to test the `create` service method. :param image_post_metadata_data: Dictionary containing the image metadata data as would be required for an `ImagePostMetadataSchema`. + :filename: Filename of the image. """ self._image_post_metadata = ImagePostMetadataSchema(**image_post_metadata_data) - self._upload_file = UploadFile(MagicMock(), size=100, filename="test.png", headers=MagicMock()) + header = {"content-type": "image/png"} + self._upload_file = UploadFile(MagicMock(), size=100, filename=filename, headers=header) self._expected_image_id = ObjectId() self.mock_object_id.return_value = self._expected_image_id @@ -151,14 +153,23 @@ class TestCreate(CreateDSL): def test_create(self): """Test creating an image.""" - self.mock_create(IMAGE_POST_METADATA_DATA_ALL_VALUES) + self.mock_create(IMAGE_POST_METADATA_DATA_ALL_VALUES, "test.png") self.call_create() self.check_create_success() + def test_create_with_invalid_file(self): + """Test creating an image with an inconsistent file extension and content type.""" + + self.mock_create(IMAGE_POST_METADATA_DATA_ALL_VALUES, "test.jpeg") + self.call_create_expecting_error(InvalidObjectIdError) + self.check_create_failed_with_exception( + f"File extension `{self._upload_file.filename}` does not match content type `{self._upload_file.content_type}`" + ) + def test_create_with_invalid_entity_id(self): """Test creating an image with an invalid `entity_id`.""" - self.mock_create({**IMAGE_POST_METADATA_DATA_ALL_VALUES, "entity_id": "invalid-id"}) + self.mock_create({**IMAGE_POST_METADATA_DATA_ALL_VALUES, "entity_id": "invalid-id"}, "test.png") self.call_create_expecting_error(InvalidObjectIdError) self.check_create_failed_with_exception("Invalid ObjectId value 'invalid-id'")