Test and perfection for cookie handling.
authorIan Cordasco <graffatcolmingov@gmail.com>
Mon, 11 Feb 2013 00:36:36 +0000 (19:36 -0500)
committerIan Cordasco <graffatcolmingov@gmail.com>
Mon, 11 Feb 2013 00:36:36 +0000 (19:36 -0500)
I also fixed up some of the RequestsCookieJar methods so using
jar.update(other_jar) works without a problem. This cleans up some of the code
in sessions and the resolve_redirects method.

requests/cookies.py
requests/sessions.py
test_requests.py

index bd7289eef29115900ad2895b793573392382a526..eb6c3146106a9b196db76d608eae61383394908a 100644 (file)
@@ -240,12 +240,18 @@ class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping):
         """Dict-like __getitem__() for compatibility with client code. Throws exception
         if there are more than one cookie with name. In that case, use the more
         explicit get() method instead. Caution: operation is O(n), not O(1)."""
+        if isinstance(name, cookielib.Cookie):
+            name = name.name
+
         return self._find_no_duplicates(name)
 
     def __setitem__(self, name, value):
         """Dict-like __setitem__ for compatibility with client code. Throws exception
         if there is already a cookie of that name in the jar. In that case, use the more
         explicit set() method instead."""
+        if isinstance(name, cookielib.Cookie):
+            name = name.name
+
         self.set(name, value)
 
     def __delitem__(self, name):
index 37b5e2054c5d9c8886be43929d2934d55b128ad0..ff70db7481a4505469ca4ca856d8016f690937ce 100644 (file)
@@ -85,6 +85,7 @@ class SessionRedirectMixin(object):
         prepared_request.hooks = req.hooks
         prepared_request.method = req.method
         prepared_request.url = req.url
+        cookiejar = resp.cookies
 
         # ((resp.status_code is codes.see_other))
         while (('location' in resp.headers and resp.status_code in REDIRECT_STATI)):
@@ -131,13 +132,13 @@ class SessionRedirectMixin(object):
 
                 prepared_request.body = None
 
+            headers = prepared_request.headers
             try:
-                del prepared_request.headers['Cookie']
+                del headers['Cookie']
             except KeyError:
                 pass
 
-            if response.headers.get('Set-Cookie'):
-                prepared_request.headers['Cookie'] = response.headers.get('Set-Cookie')
+            prepared_request.prepare_cookies(cookiejar)
 
             resp = self.send(
                 prepared_request,
@@ -149,9 +150,13 @@ class SessionRedirectMixin(object):
                 allow_redirects=False,
             )
 
+            cookiejar.update(resp.cookies)
+
             i += 1
             yield resp
 
+        resp.cookies.update(cookiejar)
+
 
 class Session(SessionRedirectMixin):
     """A Requests session.
index e1317f53c3e23e7de73d1aa7e7e2d0f15713099f..6c4c3dd21414ef675f917c49329a55f0a08c2712 100644 (file)
@@ -112,7 +112,8 @@ class RequestsTestCase(unittest.TestCase):
         self.assertEqual(r.status_code, 200)
 
     def test_set_cookie_on_301(self):
-        url = httpbin('cookies/set/foo/bar')
+        s = requests.session()
+        url = httpbin('cookies/set?foo=bar')
         r = s.get(url)
         self.assertTrue(s.cookies['foo'] == 'bar')