Mercurial > public > sg101
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) +