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 bd7289e..eb6c314 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 37b5e20..ff70db7 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 e1317f5..6c4c3dd 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')