Changing requests to use connection pooling

Updated tests and mocks to test for connection pooling
twilio · Sep 18, 2012 · 75bde43 · 75bde43
1 parent a156011
commit 75bde43
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 97 deletions.
diff --git a/src/shadow/proxy/web.py b/src/shadow/proxy/web.py
@@ -62,6 +62,12 @@ def __init__(self, service,
         self.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE'])(self.catch_all)
         self.service = service
 
+        self.session = requests.session(config={
+                'pool_connections': 20,
+                'pool_maxsize': 20,
+                'keep_alive': True,
+            })
+
     def log_result(self, result):
         logger.debug('Logging results: {}'.format(result))
         for results_log in self.result_loggers:
@@ -142,7 +148,7 @@ def catch_all(self, path):
         data = dict(request.form)
         url = request.path
 
-        true_greenlets = [self.service.spawn(self.timer, timed_func=requests.request,
+        true_greenlets = [self.service.spawn(self.timer, timed_func=self.session.request,
                 method=method,
                 url='{server}{path}'.format(server=service, path=path),
                 headers=dict(headers.items() + self.true_servers_additional_headers),
@@ -152,7 +158,7 @@ def catch_all(self, path):
                 timeout=self.true_servers_timeout
             ) for service in self.true_servers]
 
-        shadow_greenlets = [self.service.spawn(self.timer, timed_func=requests.request,
+        shadow_greenlets = [self.service.spawn(self.timer, timed_func=self.session.request,
                 method=method,
                 url='{server}{path}'.format(server=service, path=path),
                 headers=dict(headers.items() + self.shadow_servers_additional_headers),

diff --git a/tests/unit/test_proxy_flask.py b/tests/unit/test_proxy_flask.py
@@ -294,19 +294,13 @@ def test_no_content_length():
 
     svc = gevent
 
-    requests = shadow.proxy.web.requests
-    requests.request = mock.Mock()
-
-    resp, expected_resp, elapsed_time = _mock_response()
 
     req, expected_req = _mock_request()
 
     req.headers = {'header': 'header_value', 'Content-Length': '231'}
 
     shadow.proxy.web.request = req
 
-    requests.request.return_value = resp
-
     mock_result_logger = mock.Mock(spec=AbstractResultsLogger)
 
     additional_headers = [('add_header', 'header_value'), ('header', 'altered_header_value')]
@@ -332,47 +326,55 @@ def test_no_content_length():
         result_loggers=[mock_result_logger]
     )
 
-    path = "/"
-
-    # mock timer to return the randomly generated time
-    app.timer = lambda timed_func, *args, **kwargs: (timed_func(*args, **kwargs), elapsed_time)
-
-    app.catch_all(path)
-
-    requests.request.assert_has_calls([call(
-        url="true_server/",
-        headers=dict(expected_req['headers'].items() + additional_true_headers),
-        data=dict(expected_req['post'].items() + additional_true_post),
-        params=dict(expected_req['get'].items() + additional_true_get),
-        timeout=1339.0,
-        method=expected_req['method'],
-        config=shadow.proxy.web.config
-    ), call(
-        url="shadow_server/",
-        headers=dict(expected_req['headers'].items() + additional_headers),
-        data=dict(expected_req['post'].items() + additional_post),
-        params=dict(expected_req['get'].items() + additional_get),
-        timeout=1337.0,
-        method=expected_req['method'],
-        config=shadow.proxy.web.config
-    )], any_order=True)
-
-    mock_result_logger.log_result.assert_called_with({
-            'request': {
-                'modified': expected_req,
-                'original': {
-                    'url': req.path,
-                    'method': req.method,
-                    'headers': dict([(unicode(k), unicode(v)) for k, v in req.headers.items()]),
-                    'get': dict([(unicode(k), unicode(v)) for k, v in req.args.items()]),
-                    'post': dict([(unicode(k), unicode(v)) for k, v in req.form.items()])
-                }
-            },
-            'results': [
-                expected_resp,
-                expected_resp
-            ]
-        })
+    with mock.patch.object(app, 'session') as requests:
+
+        requests.request = mock.Mock()
+
+        resp, expected_resp, elapsed_time = _mock_response()
+
+        requests.request.return_value = resp
+
+        path = "/"
+
+        # mock timer to return the randomly generated time
+        app.timer = lambda timed_func, *args, **kwargs: (timed_func(*args, **kwargs), elapsed_time)
+
+        app.catch_all(path)
+
+        requests.request.assert_has_calls([call(
+            url="true_server/",
+            headers=dict(expected_req['headers'].items() + additional_true_headers),
+            data=dict(expected_req['post'].items() + additional_true_post),
+            params=dict(expected_req['get'].items() + additional_true_get),
+            timeout=1339.0,
+            method=expected_req['method'],
+            config=shadow.proxy.web.config
+        ), call(
+            url="shadow_server/",
+            headers=dict(expected_req['headers'].items() + additional_headers),
+            data=dict(expected_req['post'].items() + additional_post),
+            params=dict(expected_req['get'].items() + additional_get),
+            timeout=1337.0,
+            method=expected_req['method'],
+            config=shadow.proxy.web.config
+        )], any_order=True)
+
+        mock_result_logger.log_result.assert_called_with({
+                'request': {
+                    'modified': expected_req,
+                    'original': {
+                        'url': req.path,
+                        'method': req.method,
+                        'headers': dict([(unicode(k), unicode(v)) for k, v in req.headers.items()]),
+                        'get': dict([(unicode(k), unicode(v)) for k, v in req.args.items()]),
+                        'post': dict([(unicode(k), unicode(v)) for k, v in req.form.items()])
+                    }
+                },
+                'results': [
+                    expected_resp,
+                    expected_resp
+                ]
+            })
 
 
 
@@ -388,17 +390,10 @@ def test_catch_all_default():
 
     svc = gevent
 
-    requests = shadow.proxy.web.requests
-    requests.request = mock.Mock()
-
-    resp, expected_resp, elapsed_time = _mock_response()
-
     req, expected_req = _mock_request()
 
     shadow.proxy.web.request = req
 
-    requests.request.return_value = resp
-
     mock_result_logger = mock.Mock(spec=AbstractResultsLogger)
 
     additional_headers = [('add_header', 'header_value'), ('header', 'altered_header_value')]
@@ -424,44 +419,52 @@ def test_catch_all_default():
         result_loggers=[mock_result_logger]
     )
 
-    path = "/"
-
-    # mock timer to return the randomly generated time
-    app.timer = lambda timed_func, *args, **kwargs: (timed_func(*args, **kwargs), elapsed_time)
-
-    app.catch_all(path)
-
-    requests.request.assert_has_calls([call(
-        url="true_server/",
-        headers=dict(expected_req['headers'].items() + additional_true_headers),
-        data=dict(expected_req['post'].items() + additional_true_post),
-        params=dict(expected_req['get'].items() + additional_true_get),
-        timeout=1339.0,
-        method=expected_req['method'],
-        config=shadow.proxy.web.config
-    ), call(
-        url="shadow_server/",
-        headers=dict(expected_req['headers'].items() + additional_headers),
-        data=dict(expected_req['post'].items() + additional_post),
-        params=dict(expected_req['get'].items() + additional_get),
-        timeout=1337.0,
-        method=expected_req['method'],
-        config=shadow.proxy.web.config
-    )], any_order=True)
-
-    mock_result_logger.log_result.assert_called_with({
-            'request': {
-                'modified': expected_req,
-                'original': {
-                    'url': req.path,
-                    'method': req.method,
-                    'headers': dict([(unicode(k), unicode(v)) for k, v in req.headers.items()]),
-                    'get': dict([(unicode(k), unicode(v)) for k, v in req.args.items()]),
-                    'post': dict([(unicode(k), unicode(v)) for k, v in req.form.items()])
-                }
-            },
-            'results': [
-                expected_resp,
-                expected_resp
-            ]
-        })
+    with mock.patch.object(app, 'session') as requests:
+
+        requests.request = mock.Mock()
+
+        resp, expected_resp, elapsed_time = _mock_response()
+
+        requests.request.return_value = resp
+
+        path = "/"
+
+        # mock timer to return the randomly generated time
+        app.timer = lambda timed_func, *args, **kwargs: (timed_func(*args, **kwargs), elapsed_time)
+
+        app.catch_all(path)
+
+        requests.request.assert_has_calls([call(
+            url="true_server/",
+            headers=dict(expected_req['headers'].items() + additional_true_headers),
+            data=dict(expected_req['post'].items() + additional_true_post),
+            params=dict(expected_req['get'].items() + additional_true_get),
+            timeout=1339.0,
+            method=expected_req['method'],
+            config=shadow.proxy.web.config
+        ), call(
+            url="shadow_server/",
+            headers=dict(expected_req['headers'].items() + additional_headers),
+            data=dict(expected_req['post'].items() + additional_post),
+            params=dict(expected_req['get'].items() + additional_get),
+            timeout=1337.0,
+            method=expected_req['method'],
+            config=shadow.proxy.web.config
+        )], any_order=True)
+
+        mock_result_logger.log_result.assert_called_with({
+                'request': {
+                    'modified': expected_req,
+                    'original': {
+                        'url': req.path,
+                        'method': req.method,
+                        'headers': dict([(unicode(k), unicode(v)) for k, v in req.headers.items()]),
+                        'get': dict([(unicode(k), unicode(v)) for k, v in req.args.items()]),
+                        'post': dict([(unicode(k), unicode(v)) for k, v in req.form.items()])
+                    }
+                },
+                'results': [
+                    expected_resp,
+                    expected_resp
+                ]
+            })