bug 1247728 - Use sessions for requests

taskcluster/client.py

@@ -41,14 +41,21 @@
 }
 
 
+def createSession(*args, **kwargs):
+    """ Create a new requests session.  This passes through all positional and
+    keyword arguments to the requests.Session() constructor
+    """
+    return requests.Session(*args, **kwargs)
+
+
 class BaseClient(object):
     """ Base Class for API Client Classes. Each individual Client class
     needs to set up its own methods for REST endpoints and Topic Exchange
     routing key patterns.  The _makeApiCall() and _topicExchange() methods
     help with this.
     """
 
-    def __init__(self, options=None):
+    def __init__(self, options=None, session=None):
         o = copy.deepcopy(self.classOptions)
         o.update(_defaultConfig)
         if options:
@@ -69,6 +76,11 @@ def __init__(self, options=None):
             log.debug('credentials key scrubbed from logging output')
         log.debug(dict((k, v) for k, v in o.items() if k != 'credentials'))
 
+        if session:
+            self.session = options.session
+        else:
+            self.session = createSession()
+
     def makeHawkExt(self):
         """ Make an 'ext' for Hawk authentication """
         o = self.options

taskcluster/utils.py

@@ -138,7 +138,7 @@ def scope_match(assumed_scopes, required_scope_sets):
     return False           # none of the required_scope_sets were satisfied
 
 
-def makeHttpRequest(method, url, payload, headers, retries=MAX_RETRIES):
+def makeHttpRequest(method, url, payload, headers, retries=MAX_RETRIES, session=None):
     """ Make an HTTP request and retry it until success, return request """
     retry = -1
     response = None
@@ -156,7 +156,7 @@ def makeHttpRequest(method, url, payload, headers, retries=MAX_RETRIES):
 
         log.debug('Making attempt %d', retry)
         try:
-            response = makeSingleHttpRequest(method, url, payload, headers)
+            response = makeSingleHttpRequest(method, url, payload, headers, session)
         except requests.exceptions.RequestException as rerr:
             if retry < retries:
                 log.warn('Retrying because of: %s' % rerr)
@@ -180,12 +180,13 @@ def makeHttpRequest(method, url, payload, headers, retries=MAX_RETRIES):
     assert False, "Error from last retry should have been raised!"
 
 
-def makeSingleHttpRequest(method, url, payload, headers):
+def makeSingleHttpRequest(method, url, payload, headers, session=None):
     method = method.upper()
     log.debug('Making a %s request to %s', method, url)
     log.debug('HTTP Headers: %s' % str(headers))
     log.debug('HTTP Payload: %s (limit 100 char)' % str(payload)[:100])
-    response = requests.request(method.upper(), url, data=payload, headers=headers)
+    obj = session if session else requests
+    response = obj.request(method.upper(), url, data=payload, headers=headers)
     log.debug('Received HTTP Status:    %s' % response.status_code)
     log.debug('Received HTTP Headers: %s' % str(response.headers))
     log.debug('Received HTTP Payload: %s (limit 1024 char)' % str(response.text)[:1024])

test/test_client.py

@@ -628,6 +628,12 @@ def test_insert_to_index(self):
             }
             result = self.i.insertTask('testing', payload)
             self.assertEqual(payload['expires'], result['expires'])
+            payload['rank'] += 1
+            result = self.i.insertTask('testing', payload)
+            self.assertEqual(payload['expires'], result['expires'])
+            payload['rank'] += 1
+            result = self.i.insertTask('testing', payload)
+            self.assertEqual(payload['expires'], result['expires'])
 
         def test_listworkertypes_signed_url(self):
             surl = self.a.buildSignedUrl('listWorkerTypes')

test/test_utils.py

@@ -167,7 +167,7 @@ class TestPutfile(base.TCTest):
     def test_success_put_file(self):
         with mock.patch.object(subject, 'makeSingleHttpRequest') as p:
             subject.putFile('setup.py', 'http://www.example.com', 'text/plain')
-            p.assert_called_once_with('put', 'http://www.example.com', mock.ANY, mock.ANY)
+            p.assert_called_once_with('put', 'http://www.example.com', mock.ANY, mock.ANY, mock.ANY)
 
 
 class TestStableSlugIdClosure(TestCase):