changeset 976:f5aa74dcdd7a

Ensure temporary files get deleted during hotlinking.
author Brian Neal <bgneal@gmail.com>
date Mon, 05 Oct 2015 20:07:44 -0500 (2015-10-06)
parents 8c3d52b7cbd1
children ba0f4a5f1fca
files core/download.py core/functions.py user_photos/forms.py user_photos/tests/test_forms.py
diffstat 4 files changed, 32 insertions(+), 13 deletions(-) [+]
line wrap: on
line diff
--- a/core/download.py	Mon Oct 05 20:06:24 2015 -0500
+++ b/core/download.py	Mon Oct 05 20:07:44 2015 -0500
@@ -49,6 +49,7 @@
             shutil.copyfileobj(r.raw, fp)
     except requests.RequestException:
         logger.exception("download_file download exception")
+        os.remove(path)
         raise
 
     file_size = os.stat(path).st_size
--- a/core/functions.py	Mon Oct 05 20:06:24 2015 -0500
+++ b/core/functions.py	Mon Oct 05 20:07:44 2015 -0500
@@ -1,4 +1,5 @@
 """This file houses various core utility functions"""
+from contextlib import contextmanager
 import datetime
 import logging
 import os
@@ -31,6 +32,15 @@
         os.remove(self.filename)
 
 
+@contextmanager
+def remove_file(path):
+    """Context manager for removing a file when the context is exited."""
+    try:
+        yield path
+    finally:
+        os.remove(path)
+
+
 def send_mail(subject, message, from_email, recipient_list, reply_to=None, defer=True):
     """
     The main send email function. Use this function to send email from the
--- a/user_photos/forms.py	Mon Oct 05 20:06:24 2015 -0500
+++ b/user_photos/forms.py	Mon Oct 05 20:07:44 2015 -0500
@@ -8,7 +8,7 @@
 from django.conf import settings
 
 from core.download import download_file
-from core.functions import TemporaryFile
+from core.functions import remove_file, TemporaryFile
 from core.images.upload import upload
 from core.s3 import S3Bucket
 from core.services import get_redis_connection
@@ -150,20 +150,20 @@
             return url
 
         # Try to download the file
-        path = download_file(url)
+        with remove_file(download_file(url)) as path:
 
-        # Upload it to our S3 bucket
-        bucket = S3Bucket(access_key=settings.USER_PHOTOS_ACCESS_KEY,
-                          secret_key=settings.USER_PHOTOS_SECRET_KEY,
-                          base_url=settings.HOT_LINK_PHOTOS_BASE_URL,
-                          bucket_name=settings.HOT_LINK_PHOTOS_BUCKET)
+            # Upload it to our S3 bucket
+            bucket = S3Bucket(access_key=settings.USER_PHOTOS_ACCESS_KEY,
+                              secret_key=settings.USER_PHOTOS_SECRET_KEY,
+                              base_url=settings.HOT_LINK_PHOTOS_BASE_URL,
+                              bucket_name=settings.HOT_LINK_PHOTOS_BUCKET)
 
-        now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-        metadata = {'user': self.user.username, 'date': now}
+            now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+            metadata = {'user': self.user.username, 'date': now}
 
-        url, _ = upload(filename=path,
-                        bucket=bucket,
-                        metadata=metadata,
-                        new_size=settings.USER_PHOTOS_MAX_SIZE)
+            url, _ = upload(filename=path,
+                            bucket=bucket,
+                            metadata=metadata,
+                            new_size=settings.USER_PHOTOS_MAX_SIZE)
 
         return url
--- a/user_photos/tests/test_forms.py	Mon Oct 05 20:06:24 2015 -0500
+++ b/user_photos/tests/test_forms.py	Mon Oct 05 20:07:44 2015 -0500
@@ -1,6 +1,7 @@
 """
 Tests for the forms in the user_photos application.
 """
+from contextlib import contextmanager
 
 import mock
 from django import forms
@@ -11,6 +12,11 @@
 from user_photos.forms import HotLinkImageForm
 
 
+@contextmanager
+def fake_remove_file(path):
+    yield path
+
+
 class HotLinkImageFormTestCase(TestCase):
 
     def setUp(self):
@@ -41,6 +47,7 @@
 
     @mock.patch('user_photos.forms.rate_limit_user')
     @mock.patch('user_photos.forms.download_file')
+    @mock.patch('user_photos.forms.remove_file', new=fake_remove_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):
@@ -56,6 +63,7 @@
 
     @mock.patch('user_photos.forms.rate_limit_user')
     @mock.patch('user_photos.forms.download_file')
+    @mock.patch('user_photos.forms.remove_file', new=fake_remove_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):