From 7b6e5da84b8a93e9af37545b16a6ec2f874be97e Mon Sep 17 00:00:00 2001
From: Theodore Keloglou <zf@sirodoht.com>
Date: Tue, 5 Jan 2021 23:15:21 +0000
Subject: [PATCH] fix raw not found cases and add relevant tests

---
 main/tests/test_images.py | 47 +++++++++++++++++++++++++++++++++++++++
 main/views.py             |  4 ++--
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/main/tests/test_images.py b/main/tests/test_images.py
index 5afcfee..7dc1005 100644
--- a/main/tests/test_images.py
+++ b/main/tests/test_images.py
@@ -46,6 +46,53 @@ class ImageDetailTestCase(TestCase):
         self.assertContains(response, "Uploaded on")
 
 
+class ImageRawTestCase(TestCase):
+    def setUp(self):
+        self.user = models.User.objects.create(username="alice")
+        self.user.set_password("abcdef123456")
+        self.user.save()
+        self.client.login(username="alice", password="abcdef123456")
+        with open("main/tests/testdata/vulf.jpeg", "rb") as fp:
+            self.client.post(reverse("image_list"), {"file": fp})
+        self.image = models.Image.objects.get(name="vulf")
+
+    def test_image_raw(self):
+        response = self.client.get(
+            reverse("image_raw", args=(self.image.slug, self.image.extension)),
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(self.image.data.tobytes(), response.content)
+
+
+class ImageRawWrongExtTestCase(TestCase):
+    def setUp(self):
+        self.user = models.User.objects.create(username="alice")
+        self.user.set_password("abcdef123456")
+        self.user.save()
+        self.client.login(username="alice", password="abcdef123456")
+        with open("main/tests/testdata/vulf.jpeg", "rb") as fp:
+            self.client.post(reverse("image_list"), {"file": fp})
+        self.image = models.Image.objects.get(name="vulf")
+
+    def test_image_raw(self):
+        response = self.client.get(
+            reverse("image_raw", args=(self.image.slug, "png")),
+        )
+        self.assertEqual(response.status_code, 404)
+
+
+class ImageRawNotFoundTestCase(TestCase):
+    def setUp(self):
+        self.slug = "nonexistent-slug"
+        self.extension = "jpeg"
+
+    def test_image_raw(self):
+        response = self.client.get(
+            reverse("image_raw", args=(self.slug, self.extension)),
+        )
+        self.assertEqual(response.status_code, 404)
+
+
 class ImageUpdateTestCase(TestCase):
     def setUp(self):
         self.user = models.User.objects.create(username="alice")
diff --git a/main/views.py b/main/views.py
index 1d0b586..36e57ed 100644
--- a/main/views.py
+++ b/main/views.py
@@ -423,8 +423,8 @@ class BlogImport(LoginRequiredMixin, FormView):
 
 
 def image_raw(request, slug, extension):
-    image = models.Image.objects.get(slug=slug)
-    if extension != image.extension:
+    image = models.Image.objects.filter(slug=slug).first()
+    if not image or extension != image.extension:
         raise Http404()
     return HttpResponse(image.data, content_type="image/" + image.extension)
 
-- 
2.30.1