# HG changeset patch # User Brian Neal # Date 1443746685 18000 # Node ID 6f55c086db1eb17597b5759ace2bd1f3d245597c # Parent 7138883966b3bf118aedc1b5683b71a7321037c9 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. diff -r 7138883966b3 -r 6f55c086db1e core/download.py --- 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: diff -r 7138883966b3 -r 6f55c086db1e user_photos/tests/test_forms.py --- 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) +