port manifest-fetching code to use requests instead of urllib2 (bug 941480)

Author: cvan

mkt/developers/tasks.py

@@ -4,19 +4,18 @@
 import logging
 import os
 import sys
-import time
 import traceback
-import urllib2
 import urlparse
 import uuid
 import zipfile
+from datetime import date
 
 from django import forms
 from django.conf import settings
 from django.core.files.storage import default_storage as storage
-from django.utils import translation
 from django.utils.http import urlencode
 
+import requests
 from appvalidator import validate_app, validate_packaged_app
 from celery_tasktree import task_with_callbacks
 from celeryutils import task
@@ -28,18 +27,20 @@
 from addons.models import Addon
 from amo.decorators import set_modified_on, write
 from amo.helpers import absolutify
-from amo.utils import (remove_icons, resize_image, send_mail_jinja, strip_bom,
-                       to_language)
+from amo.utils import remove_icons, resize_image, send_mail_jinja, strip_bom
 from files.models import FileUpload, File, FileValidation
 from files.utils import SafeUnzip
 
 from mkt.constants import APP_PREVIEW_SIZES
-from mkt.constants.regions import REGIONS_CHOICES_SLUG
 from mkt.webapps.models import AddonExcludedRegion, Webapp
 
 
 log = logging.getLogger('z.mkt.developers.task')
 
+CT_URL = (
+    'https://developer.mozilla.org/docs/Web/Apps/Manifest#Serving_manifests'
+)
+
 
 @task
 @write
@@ -205,19 +206,21 @@ def get_preview_sizes(ids, **kw):
 def _fetch_content(url):
     with statsd.timer('developers.tasks.fetch_content'):
         try:
-            stream = urllib2.urlopen(url, timeout=30)
-            if not 200 <= stream.getcode() < 300:
-                raise Exception(
-                    'An invalid HTTP status code was returned.')
-            if not stream.headers.keys():
-                raise Exception(
-                    'The HTTP server did not return headers.')
-            return stream
-        except urllib2.HTTPError, e:
-            raise Exception(
-                '%s responded with %s (%s).' % (url, e.code, e.msg))
-        except urllib2.URLError, e:
-            # Unpack the URLError to try and find a useful message.
+            res = requests.get(url, timeout=30, stream=True)
+
+            if not 200 <= res.status_code < 300:
+                statsd.incr('developers.tasks.fetch_content.error')
+                raise Exception('An invalid HTTP status code was returned.')
+
+            if not res.headers.keys():
+                statsd.incr('developers.tasks.fetch_content.error')
+                raise Exception('The HTTP server did not return headers.')
+
+            statsd.incr('developers.tasks.fetch_content.success')
+            return res
+        except requests.RequestException as e:
+            statsd.incr('developers.tasks.fetch_content.error')
+            log.error('fetch_content connection error: %s' % e)
             raise Exception('The file could not be retrieved.')
 
 
@@ -228,7 +231,8 @@ class ResponseTooLargeException(Exception):
 def get_content_and_check_size(response, max_size):
     # Read one extra byte. Reject if it's too big so we don't have issues
     # downloading huge files.
-    content = response.read(max_size + 1)
+    content = response.iter_content(chunk_size=max_size + 1,
+                                    decode_unicode=True).next()
     if len(content) > max_size:
         raise ResponseTooLargeException('Too much data.')
     return content
@@ -330,8 +334,6 @@ def failed_validation(*messages, **kwargs):
                        'prelim': True})
 
 
-CT_URL = (
-    'https://developer.mozilla.org/docs/Web/Apps/Manifest#Serving_manifests')
 def _fetch_manifest(url, upload=None):
     def fail(message, upload=None):
         if upload is None:
@@ -348,7 +350,7 @@ def fail(message, upload=None):
                'again.'), upload=upload)
         return
 
-    ct = response.headers.get('Content-Type', '')
+    ct = response.headers.get('content-type', '')
     if not ct.startswith('application/x-web-app-manifest+json'):
         fail(_('Manifests must be served with the HTTP header '
                '"Content-Type: application/x-web-app-manifest+json". See %s '
@@ -416,17 +418,17 @@ def subscribe_to_responsys(campaign, address, format='html', source_url='',
         'SOURCE_URL': source_url,
         'EMAIL_ADDRESS_': address,
         'EMAIL_FORMAT_': 'H' if format == 'html' else 'T',
-        }
+    }
 
     data['%s_FLG' % campaign] = 'Y'
     data['%s_DATE' % campaign] = date.today().strftime('%Y-%m-%d')
     data['_ri_'] = settings.RESPONSYS_ID
 
     try:
-        res = urllib2.urlopen('http://awesomeness.mozilla.org/pub/rf',
-                              data=urlencode(data))
-        return res.code == 200
-    except urllib2.URLError:
+        res = requests.get('http://awesomeness.mozilla.org/pub/rf',
+                           data=urlencode(data))
+        return res.status_code == 200
+    except requests.RequestException:
         return False
 
 
@@ -457,8 +459,9 @@ def region_email(ids, regions, **kw):
         to = set(product.authors.values_list('email', flat=True))
 
         if len(regions) == 1:
-            subject = _(u'{region} region added to the Firefox Marketplace'
-                ).format(region=regions[0])
+            subject = _(
+                u'{region} region added to the Firefox Marketplace').format(
+                    region=regions[0])
         else:
             subject = _(u'New regions added to the Firefox Marketplace')
 

mkt/developers/tests/test_tasks.py

@@ -1,12 +1,12 @@
 import codecs
-from contextlib import contextmanager
-from cStringIO import StringIO
 import json
 import os
 import shutil
 import socket
 import tempfile
 import urllib2
+from contextlib import contextmanager
+from cStringIO import StringIO
 
 from django.conf import settings
 from django.core import mail
@@ -15,6 +15,7 @@
 import mock
 from nose.tools import eq_
 from PIL import Image
+from requests import RequestException
 
 import amo
 import amo.tests
@@ -168,6 +169,8 @@ def test_validate_packaged_app(self, _zipfile, _mock):
 
 
 storage_open = storage.open
+
+
 def _mock_hide_64px_icon(path, *args, **kwargs):
     """
     A function that mocks `storage.open` and throws an IOError if you try to
@@ -214,8 +217,8 @@ def setUp(self):
         self.upload = FileUpload.objects.create()
         self.content_type = 'application/x-web-app-manifest+json'
 
-        patcher = mock.patch('mkt.developers.tasks.urllib2.urlopen')
-        self.urlopen_mock = patcher.start()
+        patcher = mock.patch('mkt.developers.tasks.requests.get')
+        self.requests_mock = patcher.start()
         self.addCleanup(patcher.stop)
 
     def get_upload(self):
@@ -225,19 +228,18 @@ def file(self, name):
         return os.path.join(os.path.dirname(__file__), 'addons', name)
 
     @contextmanager
-    def patch_urlopen(self):
-        response_mock = mock.Mock()
-        response_mock.getcode.return_value = 200
-        response_mock.read.return_value = '<default>'
-        response_mock.headers = {'Content-Type': self.content_type}
+    def patch_requests(self):
+        response_mock = mock.Mock(status_code=200)
+        response_mock.iter_content.return_value = mock.Mock(
+            next=lambda: '<default>')
+        response_mock.headers = {'content-type': self.content_type}
         yield response_mock
-        self.urlopen_mock.return_value = response_mock
+        self.requests_mock.return_value = response_mock
 
     @mock.patch('mkt.developers.tasks.validator')
     def test_success_add_file(self, validator_mock):
-        with self.patch_urlopen() as ur:
-            ur.read.return_value = 'woo'
-            ur.headers = {'Content-Type': self.content_type}
+        with self.patch_requests() as ur:
+            ur.iter_content.return_value = mock.Mock(next=lambda: 'woo')
 
         tasks.fetch_manifest('http://xx.com/manifest.json', self.upload.pk)
         upload = FileUpload.objects.get(pk=self.upload.pk)
@@ -247,9 +249,9 @@ def test_success_add_file(self, validator_mock):
 
     @mock.patch('mkt.developers.tasks.validator')
     def test_success_call_validator(self, validator_mock):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             ct = self.content_type + '; charset=utf-8'
-            ur.headers = {'Content-Type': ct}
+            ur.headers = {'content-type': ct}
 
         tasks.fetch_manifest('http://xx.com/manifest.json', self.upload.pk)
         assert validator_mock.called
@@ -273,31 +275,31 @@ def check_validation(self, msg=''):
 
     def test_connection_error(self):
         reason = socket.gaierror(8, 'nodename nor servname provided')
-        self.urlopen_mock.side_effect = urllib2.URLError(reason)
+        self.requests_mock.side_effect = RequestException(reason)
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation(
             'No manifest was found at that URL. Check the address and try '
             'again.')
 
     def test_url_timeout(self):
         reason = socket.timeout('too slow')
-        self.urlopen_mock.side_effect = urllib2.URLError(reason)
+        self.requests_mock.side_effect = RequestException(reason)
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation(
             'No manifest was found at that URL. Check the address and try '
             'again.')
 
     def test_other_url_error(self):
         reason = Exception('Some other failure.')
-        self.urlopen_mock.side_effect = urllib2.URLError(reason)
+        self.requests_mock.side_effect = RequestException(reason)
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation(
             'No manifest was found at that URL. Check the address and try '
             'again.')
 
     @mock.patch('mkt.developers.tasks.validator', lambda uid, **kw: None)
     def test_no_content_type(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             ur.headers = {}
 
         tasks.fetch_manifest('url', self.upload.pk)
@@ -307,7 +309,7 @@ def test_no_content_type(self):
 
     @mock.patch('mkt.developers.tasks.validator', lambda uid, **kw: None)
     def test_bad_content_type(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             ur.headers = {'Content-Type': 'x'}
 
         tasks.fetch_manifest('url', self.upload.pk)
@@ -318,63 +320,67 @@ def test_bad_content_type(self):
 
     @mock.patch('mkt.developers.tasks.validator', lambda uid, **kw: None)
     def test_good_charset(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             ur.headers = {
-                'Content-Type': 'application/x-web-app-manifest+json;'
-                                'charset=utf-8'}
+                'content-type': 'application/x-web-app-manifest+json;'
+                                'charset=utf-8'
+            }
 
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation()
 
     @mock.patch('mkt.developers.tasks.validator', lambda uid, **kw: None)
     def test_bad_charset(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             ur.headers = {
-                'Content-Type': 'application/x-web-app-manifest+json;'
-                                'charset=ISO-1234567890-LOL'}
+                'content-type': 'application/x-web-app-manifest+json;'
+                                'charset=ISO-1234567890-LOL'
+            }
 
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation("The manifest's encoding does not match the "
                               'charset provided in the HTTP Content-Type.')
 
     def test_response_too_large(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             content = 'x' * (settings.MAX_WEBAPP_UPLOAD_SIZE + 1)
-            ur.read.return_value = content
+            ur.iter_content.return_value = mock.Mock(next=lambda: content)
 
         tasks.fetch_manifest('url', self.upload.pk)
         max_webapp_size = settings.MAX_WEBAPP_UPLOAD_SIZE
         self.check_validation('Your manifest must be less than %s bytes.' %
                               max_webapp_size)
 
     def test_http_error(self):
-        self.urlopen_mock.side_effect = urllib2.HTTPError(
+        self.requests_mock.side_effect = urllib2.HTTPError(
             'url', 404, 'Not Found', [], None)
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation(
             'No manifest was found at that URL. Check the address and try '
             'again.')
 
     def test_strip_utf8_bom(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             with open(self.file('utf8bom.webapp')) as fp:
-                ur.read.return_value = fp.read()
+                content = fp.read()
+                ur.iter_content.return_value = mock.Mock(next=lambda: content)
 
         tasks.fetch_manifest('url', self.upload.pk)
         upload = self.get_upload()
         with storage.open(upload.path, 'rb') as fp:
             manifest = fp.read()
-            json.loads(manifest)  # no parse error
+            json.loads(manifest)  # No parse error.
             assert not manifest.startswith(codecs.BOM_UTF8)
 
     def test_non_utf8_encoding(self):
-        with self.patch_urlopen() as ur:
+        with self.patch_requests() as ur:
             with open(self.file('utf8bom.webapp')) as fp:
-                # Set encoding to utf16 which will be invalid
-                ur.read.return_value = fp.read().decode('utf8').encode('utf16')
+                # Set encoding to utf16 which will be invalid.
+                content = fp.read().decode('utf8').encode('utf16')
+                ur.iter_content.return_value = mock.Mock(next=lambda: content)
         tasks.fetch_manifest('url', self.upload.pk)
         self.check_validation(
-                    'Your manifest file was not encoded as valid UTF-8.')
+            'Your manifest file was not encoded as valid UTF-8.')
 
 
 class TestFetchIcon(BaseWebAppTest):
@@ -384,9 +390,9 @@ def setUp(self):
         self.content_type = 'image/png'
         self.apps_path = os.path.join(settings.ROOT, 'apps', 'devhub', 'tests',
                                       'addons')
-        patcher = mock.patch('mkt.developers.tasks.urllib2.urlopen')
-        self.urlopen_mock = patcher.start()
-        self.urlopen_mock.return_value = StringIO('mozballin')
+        patcher = mock.patch('mkt.developers.tasks.requests.get')
+        self.requests_mock = patcher.start()
+        self.requests_mock.return_value = StringIO('mozballin')
         self.addCleanup(patcher.stop)
 
     def webapp_from_path(self, path):
@@ -404,13 +410,13 @@ def test_no_icons(self):
         path = os.path.join(self.apps_path, 'noicon.webapp')
         iconless_app = self.webapp_from_path(path)
         tasks.fetch_icon(iconless_app)
-        assert not self.urlopen_mock.called
+        assert not self.requests_mock.called
 
     def test_bad_icons(self):
         path = os.path.join(self.apps_path, 'badicon.webapp')
         iconless_app = self.webapp_from_path(path)
         tasks.fetch_icon(iconless_app)
-        assert not self.urlopen_mock.called
+        assert not self.requests_mock.called
 
     def check_icons(self, webapp):
         manifest = webapp.get_manifest_json()