From 615d3f732dfa206cd999547696640502139ce5dc Mon Sep 17 00:00:00 2001 From: openhands Date: Thu, 6 Feb 2025 15:35:24 +0000 Subject: [PATCH] Fix issue #194: [Claude] Refactor picture_creation --- photo/mutations.py | 76 +++++++++++++++++----------- photo/tests/test_picture_creation.py | 68 +++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 30 deletions(-) create mode 100644 photo/tests/test_picture_creation.py diff --git a/photo/mutations.py b/photo/mutations.py index d08b96e..db33930 100644 --- a/photo/mutations.py +++ b/photo/mutations.py @@ -85,39 +85,55 @@ def __init__(self, message) -> None: self.message = message +def _get_user(user_id: str) -> User: + """Get user by ID or raise ValidationError if not found.""" + if not (user := User.objects.filter(id=user_id).first()): + raise ValidationError(message=NO_USER_FOUND) + return user + +def _process_image_file(file, max_size: int) -> tuple[BytesIO, str | None]: + """Process image file, convert to webp format and validate size.""" + image_bytes = BytesIO() + filename = None + + if isinstance(file, InMemoryUploadedFile): + if file.size > max_size: + raise ValidationError(message=PICTURE_SIZE_ERROR) + image = Image.open(file) + filename = str(file.name).rsplit(".", 1)[0] + image.save(image_bytes, format="webp", optimize=True) + else: + file.save(image_bytes, format="webp") + if image_bytes.tell() > max_size: + raise ValidationError(message=PICTURE_SIZE_ERROR) + + image_bytes.seek(0) + return image_bytes, filename + +def _create_picture_object(user: User, image_bytes: BytesIO, filename: str | None) -> Picture: + """Create and save Picture object with the processed image.""" + picture_object = Picture(user=user) + picture_object.save() + + image_file = SimpleUploadedFile( + str(picture_object.id), + image_bytes.getvalue(), + content_type="image/webp", + ) + + picture_object.file = image_file + picture_object.name = filename if filename else picture_object.id + picture_object.save() + + return picture_object + def picture_creation(input: PictureInput) -> CreatePictureMutationResponse: + """Create a new picture from the provided input.""" try: with transaction.atomic(): - if not (user := User.objects.filter(id=input.user).first()): - raise ValidationError(message=NO_USER_FOUND) - - picture_object = Picture(user=user) - picture_object.save() - - image_bytes = BytesIO() - file = input.file - filename = None - if type(input.file) == InMemoryUploadedFile: - if input.file.size > int(settings.MAX_PICTURE_SIZE): - raise ValidationError(message=PICTURE_SIZE_ERROR) - image = Image.open(input.file) - filename = str(input.file.name).rsplit(".", 1)[0] - image.save(image_bytes, format="webp", optimize=True) - else: - file.save(image_bytes, format="webp") - if image_bytes.tell() > int(settings.MAX_PICTURE_SIZE): - raise ValidationError(message=PICTURE_SIZE_ERROR) - image_bytes.seek(0) - - image_file = SimpleUploadedFile( - str(picture_object.id), - image_bytes.getvalue(), - content_type="image/webp", - ) - - picture_object.file = image_file - picture_object.name = filename if filename else picture_object.id - picture_object.save() + user = _get_user(input.user) + image_bytes, filename = _process_image_file(input.file, int(settings.MAX_PICTURE_SIZE)) + picture_object = _create_picture_object(user, image_bytes, filename) return CreatePictureMutationResponse( success=True, results=picture_object, errors="" diff --git a/photo/tests/test_picture_creation.py b/photo/tests/test_picture_creation.py new file mode 100644 index 0000000..9f54ce5 --- /dev/null +++ b/photo/tests/test_picture_creation.py @@ -0,0 +1,68 @@ +import pytest +from django.core.files.uploadedfile import InMemoryUploadedFile, SimpleUploadedFile +from django.forms import ValidationError +from PIL import Image +from io import BytesIO + +from photo.mutations import picture_creation +from photo.models import User, Picture +from photo.inputs import PictureInput +from photo.fixtures import NO_USER_FOUND, PICTURE_SIZE_ERROR, CREATE_PICTURE_ERROR + +@pytest.fixture +def test_user(): + return User.objects.create(id="test_user") + +@pytest.fixture +def test_image(): + image = Image.new('RGB', (100, 100), color='red') + image_bytes = BytesIO() + image.save(image_bytes, format='webp') + image_bytes.seek(0) + return image_bytes + +def test_picture_creation_success(test_user, test_image): + # Create test input + file = SimpleUploadedFile("test.webp", test_image.getvalue(), content_type="image/webp") + input = PictureInput(user=test_user.id, file=file) + + # Test picture creation + response = picture_creation(input=input) + + assert response.success is True + assert response.errors == "" + assert isinstance(response.results, Picture) + assert response.results.user == test_user + assert response.results.name == str(response.results.id) + +def test_picture_creation_invalid_user(): + # Create test input with non-existent user + file = SimpleUploadedFile("test.webp", b"test", content_type="image/webp") + input = PictureInput(user="non_existent_user", file=file) + + # Test picture creation + response = picture_creation(input=input) + + assert response.success is False + assert response.errors == NO_USER_FOUND + assert response.results == {} + +def test_picture_creation_large_file(test_user, settings): + # Create a large test file + large_image = Image.new('RGB', (1000, 1000), color='red') + image_bytes = BytesIO() + large_image.save(image_bytes, format='webp') + image_bytes.seek(0) + + # Set a small max file size for testing + settings.MAX_PICTURE_SIZE = 100 + + file = SimpleUploadedFile("test.webp", image_bytes.getvalue(), content_type="image/webp") + input = PictureInput(user=test_user.id, file=file) + + # Test picture creation + response = picture_creation(input=input) + + assert response.success is False + assert response.errors == PICTURE_SIZE_ERROR + assert response.results == {}