From 2e3169615636ceffc7aa85469464470c11d3a85e Mon Sep 17 00:00:00 2001
From: Ian Cordasco <graffatcolmingov@gmail.com>
Date: Sun, 10 Feb 2013 19:36:36 -0500
Subject: [PATCH] Test and perfection for cookie handling.

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  |  6 ++++++
 requests/sessions.py | 11 ++++++++---
 test_requests.py     |  3 ++-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/requests/cookies.py b/requests/cookies.py
index bd7289e..eb6c314 100644
--- a/requests/cookies.py
+++ b/requests/cookies.py
@@ -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):
diff --git a/requests/sessions.py b/requests/sessions.py
index 37b5e20..ff70db7 100644
--- a/requests/sessions.py
+++ b/requests/sessions.py
@@ -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.
diff --git a/test_requests.py b/test_requests.py
index e1317f5..6c4c3dd 100644
--- a/test_requests.py
+++ b/test_requests.py
@@ -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')
 
-- 
2.7.4