Merge pull request #63 from sigmavirus24/headers-refactor

Headers handling refactor
betamaxpy · Apr 18, 2015 · a831f3d · a831f3d
2 parents 9f0375e + c5dae8d
commit a831f3d
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 14 deletions.
diff --git a/betamax/cassette/mock_response.py b/betamax/cassette/mock_response.py
@@ -6,7 +6,7 @@ class MockHTTPResponse(object):
     def __init__(self, headers):
         from .util import coerce_content
 
-        h = ["%s: %s" % (k, v) for (k, v) in headers.items()]
+        h = ["%s: %s" % (k, v) for k in headers for v in headers.getlist(k)]
         h = map(coerce_content, h)
         h = '\r\n'.join(h)
         if sys.version_info < (2, 7):

diff --git a/betamax/cassette/util.py b/betamax/cassette/util.py
@@ -2,6 +2,7 @@
 from datetime import datetime
 from requests.models import PreparedRequest, Response
 from requests.packages.urllib3 import HTTPResponse
+from requests.packages.urllib3._collections import HTTPHeaderDict
 from requests.structures import CaseInsensitiveDict
 from requests.status_codes import _codes
 from requests.cookies import RequestsCookieJar
@@ -90,10 +91,14 @@ def deserialize_prepared_request(serialized):
 def serialize_response(response, preserve_exact_body_bytes):
     body = {'encoding': response.encoding}
     add_body(response, preserve_exact_body_bytes, body)
+    header_map = response.raw.headers
+    headers = {}
+    for header_name in header_map.keys():
+        headers[header_name] = header_map.getlist(header_name)
 
     return {
         'body': body,
-        'headers': dict((k, [v]) for k, v in response.headers.items()),
+        'headers': headers,
         'status': {'code': response.status_code, 'message': response.reason},
         'url': response.url,
     }
@@ -102,20 +107,28 @@ def serialize_response(response, preserve_exact_body_bytes):
 def deserialize_response(serialized):
     r = Response()
     r.encoding = serialized['body']['encoding']
-    h = [(k, from_list(v)) for k, v in serialized['headers'].items()]
-    r.headers = CaseInsensitiveDict(h)
+    header_dict = HTTPHeaderDict()
+
+    for header_name, header_list in serialized['headers'].items():
+        if isinstance(header_list, list):
+            for header_value in header_list:
+                header_dict.add(header_name, header_value)
+        else:
+            header_dict.add(header_name, header_list)
+    r.headers = CaseInsensitiveDict(header_dict)
+
     r.url = serialized.get('url', '')
     if 'status' in serialized:
         r.status_code = serialized['status']['code']
         r.reason = serialized['status']['message']
     else:
         r.status_code = serialized['status_code']
         r.reason = _codes[r.status_code][0].upper()
-    add_urllib3_response(serialized, r)
+    add_urllib3_response(serialized, r, header_dict)
     return r
 
 
-def add_urllib3_response(serialized, response):
+def add_urllib3_response(serialized, response, headers):
     if 'base64_string' in serialized['body']:
         body = io.BytesIO(
             base64.b64decode(serialized['body']['base64_string'].encode())
@@ -126,9 +139,9 @@ def add_urllib3_response(serialized, response):
     h = HTTPResponse(
         body,
         status=response.status_code,
-        headers=response.headers,
+        headers=headers,
         preload_content=False,
-        original_response=MockHTTPResponse(response.headers)
+        original_response=MockHTTPResponse(headers)
     )
     response.raw = h
 

diff --git a/tests/cassettes/test-multiple-cookies-regression.json b/tests/cassettes/test-multiple-cookies-regression.json
@@ -0,0 +1 @@
+{"http_interactions": [{"request": {"body": {"encoding": "utf-8", "string": ""}, "method": "GET", "headers": {"Connection": ["keep-alive"], "Accept-Encoding": ["gzip, deflate"], "User-Agent": ["python-requests/2.6.0 CPython/3.4.2 Darwin/14.1.0"], "Accept": ["*/*"]}, "uri": "https://httpbin.org/cookies/set?cookie2=value2&cookie1=value1&cookie3=value3&cookie0=value0"}, "recorded_at": "2015-04-18T16:05:26", "response": {"body": {"encoding": "utf-8", "string": "<!DOCTYPE HTML PUBLIC \"-//W3C//DTD HTML 3.2 Final//EN\">\n<title>Redirecting...</title>\n<h1>Redirecting...</h1>\n<p>You should be redirected automatically to target URL: <a href=\"/cookies\">/cookies</a>.  If not click the link."}, "url": "https://httpbin.org/cookies/set?cookie2=value2&cookie1=value1&cookie3=value3&cookie0=value0", "headers": {"date": ["Sat, 18 Apr 2015 16:05:26 GMT"], "content-length": ["223"], "set-cookie": ["cookie1=value1; Path=/", "cookie0=value0; Path=/", "cookie3=value3; Path=/", "cookie2=value2; Path=/"], "location": ["/cookies"], "connection": ["keep-alive"], "access-control-allow-origin": ["*"], "content-type": ["text/html; charset=utf-8"], "access-control-allow-credentials": ["true"], "server": ["nginx"]}, "status": {"message": "FOUND", "code": 302}}}, {"request": {"body": {"encoding": "utf-8", "string": ""}, "method": "GET", "headers": {"Connection": ["keep-alive"], "Accept-Encoding": ["gzip, deflate"], "User-Agent": ["python-requests/2.6.0 CPython/3.4.2 Darwin/14.1.0"], "Accept": ["*/*"], "Cookie": ["cookie2=value2; cookie1=value1; cookie3=value3; cookie0=value0"]}, "uri": "https://httpbin.org/cookies"}, "recorded_at": "2015-04-18T16:05:26", "response": {"body": {"encoding": null, "string": "{\n  \"cookies\": {\n    \"cookie0\": \"value0\", \n    \"cookie1\": \"value1\", \n    \"cookie2\": \"value2\", \n    \"cookie3\": \"value3\"\n  }\n}\n"}, "url": "https://httpbin.org/cookies", "headers": {"date": ["Sat, 18 Apr 2015 16:05:26 GMT"], "content-length": ["125"], "access-control-allow-credentials": ["true"], "connection": ["keep-alive"], "access-control-allow-origin": ["*"], "content-type": ["application/json"], "server": ["nginx"]}, "status": {"message": "OK", "code": 200}}}], "recorded_with": "betamax/0.4.1"}
diff --git a/tests/integration/test_multiple_cookies.py b/tests/integration/test_multiple_cookies.py
@@ -0,0 +1,42 @@
+import betamax
+
+from .helper import IntegrationHelper
+
+
+class TestMultipleCookies(IntegrationHelper):
+    """Previously our handling of multiple instances of cookies was wrong.
+
+    This set of tests is here to ensure that we properly serialize/deserialize
+    the case where the client receives and betamax serializes multiple
+    Set-Cookie headers.
+
+    See the following for more information:
+
+    - https://github.com/sigmavirus24/betamax/pull/60
+    - https://github.com/sigmavirus24/betamax/pull/59
+    - https://github.com/sigmavirus24/betamax/issues/58
+    """
+    def setUp(self):
+        super(TestMultipleCookies, self).setUp()
+        self.cassette_created = False
+
+    def test_multiple_cookies(self):
+        """Make a request to httpbin.org and verify we serialize it correctly.
+
+        We should be able to see that the cookiejar on the session has the
+        cookies properly parsed and distinguished.
+        """
+        recorder = betamax.Betamax(self.session)
+        cassette_name = 'test-multiple-cookies-regression'
+        url = 'https://httpbin.org/cookies/set'
+        cookies = {
+            'cookie0': 'value0',
+            'cookie1': 'value1',
+            'cookie2': 'value2',
+            'cookie3': 'value3',
+        }
+        with recorder.use_cassette(cassette_name):
+            self.session.get(url, params=cookies)
+
+        for name, value in cookies.items():
+            assert self.session.cookies[name] == value
diff --git a/tests/unit/test_cassette.py b/tests/unit/test_cassette.py
@@ -9,6 +9,7 @@
 from betamax.cassette import util
 from requests.models import Response, Request
 from requests.packages import urllib3
+from requests.packages.urllib3._collections import HTTPHeaderDict
 from requests.structures import CaseInsensitiveDict
 
 
@@ -63,7 +64,7 @@ def test_serialize_response(self):
                 'string': decode('foo'),
                 'encoding': 'utf-8'
             }
-        }, r)
+        }, r, HTTPHeaderDict())
         serialized = util.serialize_response(r, False)
         assert serialized is not None
         assert serialized != {}
@@ -169,7 +170,7 @@ def test_add_urllib3_response(self):
                 'string': decode('foo'),
                 'encoding': 'utf-8'
             }
-        }, r)
+        }, r, HTTPHeaderDict())
         assert isinstance(r.raw, urllib3.response.HTTPResponse)
         assert r.content == b'foo'
         assert isinstance(r.raw._original_response, cassette.MockHTTPResponse)
@@ -202,7 +203,7 @@ def setUp(self):
                 'string': decode('foo'),
                 'encoding': 'utf-8'
             }
-        }, r)
+        }, r, HTTPHeaderDict({'Content-Type': decode('foo')}))
         self.response = r
 
         # Create an associated request
@@ -236,7 +237,7 @@ def setUp(self):
                     'string': decode('foo'),
                     'encoding': 'utf-8',
                 },
-                'headers': {'Content-Type': [decode('foo')]},
+                'headers': {'content-type': [decode('foo')]},
                 'status': {'code': 200, 'message': 'OK'},
                 'url': 'http://example.com',
             },
@@ -408,9 +409,9 @@ def test_replace_in_uri(self):
 
 class TestMockHTTPResponse(unittest.TestCase):
     def setUp(self):
-        self.resp = cassette.MockHTTPResponse({
+        self.resp = cassette.MockHTTPResponse(HTTPHeaderDict({
             decode('Header'): decode('value')
-        })
+        }))
 
     def test_isclosed(self):
         assert self.resp.isclosed() is False