changeset 973:6f55c086db1e

Guess file extension based on content-type. When downloading a file, and no path is supplied to store it, guess the file extension using mimetypes and the content-type header. Also supply a unit test for the HotLinkImageForm.
author Brian Neal <bgneal@gmail.com>
date Thu, 01 Oct 2015 19:44:45 -0500
parents 7138883966b3
children d260aef91ad7
files core/download.py user_photos/tests/test_forms.py
diffstat 2 files changed, 77 insertions(+), 1 deletions(-) [+]
line wrap: on
line diff
--- a/core/download.py	Wed Sep 23 21:26:09 2015 -0500
+++ b/core/download.py	Thu Oct 01 19:44:45 2015 -0500
@@ -1,6 +1,7 @@
 """This module contains routines for downloading files."""
 
 import logging
+import mimetypes
 import os
 import shutil
 import tempfile
@@ -35,7 +36,11 @@
     # Save file data
 
     if not path:
-        fd, path = tempfile.mkstemp()
+        content_type = r.headers.get('content-type')
+        suffix = mimetypes.guess_extension(content_type) if content_type else ''
+        if suffix == '.jpe':
+            suffix = '.jpg'
+        fd, path = tempfile.mkstemp(suffix=suffix)
         os.close(fd)
 
     try:
--- a/user_photos/tests/test_forms.py	Wed Sep 23 21:26:09 2015 -0500
+++ b/user_photos/tests/test_forms.py	Thu Oct 01 19:44:45 2015 -0500
@@ -0,0 +1,71 @@
+"""
+Tests for the forms in the user_photos application.
+"""
+
+import mock
+from django import forms
+from django.conf import settings
+from django.contrib.auth.models import User
+from django.test import TestCase
+
+from user_photos.forms import HotLinkImageForm
+
+
+class HotLinkImageFormTestCase(TestCase):
+
+    def setUp(self):
+        self.username = 'test_user'
+        self.pw = 'password'
+        self.user = User.objects.create_user(self.username, '', self.pw)
+        self.user.save()
+        self.client.login(username=self.username, password=self.pw)
+
+    @mock.patch('user_photos.forms.rate_limit_user')
+    def test_no_url(self, rate_limit_mock):
+        args = {}
+        form = HotLinkImageForm(args, user=self.user)
+        self.assertFalse(form.is_valid())
+
+    @mock.patch('user_photos.forms.rate_limit_user')
+    def test_bad_url(self, rate_limit_mock):
+        args = {'url': 'jkdal;jkkls;$JSx49'}
+        form = HotLinkImageForm(args, user=self.user)
+        self.assertFalse(form.is_valid())
+
+    @mock.patch('user_photos.forms.rate_limit_user')
+    def test_rate_limit(self, rate_limit_mock):
+        rate_limit_mock.side_effect = forms.ValidationError("Rate limit exceeded")
+        args = {'url': 'http://example.com/a.jpg'}
+        form = HotLinkImageForm(args, user=self.user)
+        self.assertFalse(form.is_valid())
+
+    @mock.patch('user_photos.forms.rate_limit_user')
+    @mock.patch('user_photos.forms.download_file')
+    @mock.patch('user_photos.forms.S3Bucket')
+    @mock.patch('user_photos.forms.upload')
+    def test_white_listed_url(self, upload_mock, bucket_mock, dl_mock, rate_limit_mock):
+        url = 'https://{}/a.jpg'.format(settings.USER_IMAGES_SOURCES[0])
+        args = {'url': url}
+        form = HotLinkImageForm(args, user=self.user)
+        self.assertTrue(form.is_valid())
+        result = form.save()
+        self.assertEqual(result, url)
+        self.assertFalse(dl_mock.called)
+        self.assertFalse(bucket_mock.called)
+        self.assertFalse(upload_mock.called)
+
+    @mock.patch('user_photos.forms.rate_limit_user')
+    @mock.patch('user_photos.forms.download_file')
+    @mock.patch('user_photos.forms.S3Bucket')
+    @mock.patch('user_photos.forms.upload')
+    def test_happy_path(self, upload_mock, bucket_mock, dl_mock, rate_limit_mock):
+        url = 'http://example.com/a.jpg'
+        args = {'url': url}
+        form = HotLinkImageForm(args, user=self.user)
+        self.assertTrue(form.is_valid())
+        new_url = 'https://img.example.com/a.jpg'
+        upload_mock.return_value = (new_url, None)
+        result = form.save()
+        dl_mock.assert_called_once_with(url)
+        self.assertEqual(result, new_url)
+