Fix issue #194: [Claude] Refactor picture_creation

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=""

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 == {}