From 5018aa1994bf2be26596e46da4b70d1d835bba61 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karlhobley10@gmail.com>
Date: Wed, 1 Jul 2015 16:14:41 +0100
Subject: [PATCH] Set file size on upload

---
 .../wagtailimages/tests/test_admin_views.py   | 23 +++++++++++++++++++
 wagtail/wagtailimages/views/images.py         |  7 ++++++
 wagtail/wagtailimages/views/multiple.py       |  1 +
 3 files changed, 31 insertions(+)

diff --git a/wagtail/wagtailimages/tests/test_admin_views.py b/wagtail/wagtailimages/tests/test_admin_views.py
index 8a4b954bb..cd2f9aae1 100644
--- a/wagtail/wagtailimages/tests/test_admin_views.py
+++ b/wagtail/wagtailimages/tests/test_admin_views.py
@@ -86,6 +86,9 @@ class TestImageAddView(TestCase, WagtailTestUtils):
         self.assertEqual(image.width, 640)
         self.assertEqual(image.height, 480)
 
+        # Test that the file_size field was set
+        self.assertTrue(image.file_size)
+
     def test_add_no_file_selected(self):
         response = self.post({
             'title': "Test image",
@@ -151,6 +154,25 @@ class TestImageEditView(TestCase, WagtailTestUtils):
         image = Image.objects.get(id=self.image.id)
         self.assertEqual(image.title, "Edited")
 
+    def test_edit_with_new_image_file(self):
+        file_content = get_test_image_file().file.getvalue()
+
+        # Change the file size of the image
+        self.image.file_size = 100000
+        self.image.save()
+
+        response = self.post({
+            'title': "Edited",
+            'file': SimpleUploadedFile('new.png', file_content),
+        })
+
+        # Should redirect back to index
+        self.assertRedirects(response, reverse('wagtailimages_index'))
+
+        # Check that the image file size changed (assume it changed to the correct value)
+        image = Image.objects.get(id=self.image.id)
+        self.assertNotEqual(image.file_size, 100000)
+
     def test_with_missing_image_file(self):
         self.image.file.delete(False)
 
@@ -330,6 +352,7 @@ class TestMultipleImageUploader(TestCase, WagtailTestUtils):
         # Check image
         self.assertIn('image', response.context)
         self.assertEqual(response.context['image'].title, 'test.png')
+        self.assertTrue(response.context['image'].file_size)
 
         # Check form
         self.assertIn('form', response.context)
diff --git a/wagtail/wagtailimages/views/images.py b/wagtail/wagtailimages/views/images.py
index 81912a96b..5ae6dbb00 100644
--- a/wagtail/wagtailimages/views/images.py
+++ b/wagtail/wagtailimages/views/images.py
@@ -97,6 +97,10 @@ def edit(request, image_id):
                 # which definitely isn't what we want...
                 original_file.storage.delete(original_file.name)
                 image.renditions.all().delete()
+
+                # Set new image file size
+                image.file_size = image.file.size
+
             form.save()
 
             # Reindex the image to make sure all tags are indexed
@@ -238,6 +242,9 @@ def add(request):
         image = ImageModel(uploaded_by_user=request.user)
         form = ImageForm(request.POST, request.FILES, instance=image)
         if form.is_valid():
+            # Set image file size
+            image.file_size = image.file.size
+
             form.save()
 
             # Reindex the image to make sure all tags are indexed
diff --git a/wagtail/wagtailimages/views/multiple.py b/wagtail/wagtailimages/views/multiple.py
index 13142aa44..81540cec4 100644
--- a/wagtail/wagtailimages/views/multiple.py
+++ b/wagtail/wagtailimages/views/multiple.py
@@ -62,6 +62,7 @@ def add(request):
             # Save it
             image = form.save(commit=False)
             image.uploaded_by_user = request.user
+            image.file_size = image.file.size
             image.save()
 
             # Success! Send back an edit form for this image to the user