Merge pull request #145 from KeepSafe/cookies_share_fixes

share response cookies as soon as received (relates to #139)

Author: popravich

aiohttp/client.py

@@ -629,7 +629,7 @@ def start(self, connection, read_until_eof=False):
                 except http.cookies.CookieError as exc:
                     client_log.warning(
                         'Can not load response cookies: %s', exc)
-
+            connection.share_cookies(self.cookies)
         return self
 
     def close(self, force=False):

aiohttp/connector.py

@@ -46,6 +46,10 @@ def release(self):
             self._transport = None
             self._wr = None
 
+    def share_cookies(self, cookies):
+        if self._connector._share_cookies:  # XXX
+            self._connector.update_cookies(cookies)
+
 
 class BaseConnector(object):
     """Base connector class.
@@ -178,8 +182,6 @@ def _release(self, key, req, transport, protocol):
                 should_close = True
             else:
                 should_close = resp.message.should_close
-                if self._share_cookies and resp.cookies:
-                    self.update_cookies(resp.cookies.items())
 
         if self._force_close:
             should_close = True

docs/client.rst

@@ -452,44 +452,6 @@ requests:
 .. note::
    By default ``share_cookies`` is set to ``False``.
 
-.. _issue-139: https://github.com/KeepSafe/aiohttp/issues/139
-
-.. note::
-   Sharing cookies between requests must be clarified a bit.
-   You might get into a `situation <issue-139_>`_ when it seems like cookies are
-   not being shared between requests::
-
-      >>> conn = aiohttp.connector.TCPConnector(share_cookies=True)
-      >>> # we assume example.com sets cookies on GET without redirect
-      >>> resp1 = yield from aiohttp.request('get', 'http://example.com/',
-      ...                                    connector=conn)
-      >>> assert 'set-cookie' in resp1.headers  # assume we got cookies in response
-      >>> assert dict(conn.cookies) == {}       # but no cookies to share!
-      >>> resp2 = yield from aiohttp.request('get', 'http://example.com/',
-      ...                                    connector=conn)
-      >>> assert 'set-cookie' in resp2.headers  # cookies again!
-
-   This happens because underlying connection handled by ``TCPConnector`` still
-   belongs to ``resp1`` (because you might want to do something with response's body).
-   As soon as you release that connection by ``yield'ing from`` ``resp.read()``,
-   ``resp.text()``, ``resp.json()`` or ``resp.release()`` the connection
-   will become available (with cookies) for sebsequent requests::
-
-      >>> conn = aiohttp.connection.TCPConnector(share_cookies=True)
-      >>> # assume the same
-      >>> resp1 = yield from aiohttp.request('get', 'http://example.com/',
-      ...                                    connector=conn)
-      >>> assert 'set-cookie' in resp1.headers
-      >>> assert dict(conn.cookies) == {}
-      >>> # process response
-      >>> do_something_with_response_or_release(resp1)
-      >>> assert dict(conn.cookies) != {}
-      >>> #
-      >>> resp2 = yield from aiohttp.request('get', 'http://example.com/',
-      ...                                    connector=conn)
-      >>> assert 'set-cookie' not in resp2.headers
-      >>> assert dict(conn.cookies) != {}
-
 
 Timeouts
 --------

tests/test_client_functional.py

@@ -722,6 +722,25 @@ def test_set_cookies(self, m_log):
         m_log.warning.assert_called_with('Can not load response cookies: %s',
                                          mock.ANY)
 
+    def test_share_cookies(self):
+        with test_utils.run_server(self.loop, router=Functional) as httpd:
+            conn = aiohttp.TCPConnector(share_cookies=True, loop=self.loop)
+            resp = self.loop.run_until_complete(
+                client.request('get', httpd.url('cookies'),
+                               connector=conn, loop=self.loop))
+            self.assertIn('SET-COOKIE', resp.headers)
+            self.assertEqual(resp.cookies['c1'].value, 'cookie1')
+            self.assertEqual(resp.cookies['c2'].value, 'cookie2')
+            self.assertEqual(conn.cookies, resp.cookies)
+
+            resp2 = self.loop.run_until_complete(
+                client.request('get', httpd.url('method', 'get'),
+                               connector=conn, loop=self.loop))
+            self.assertNotIn('SET-COOKIE', resp2.headers)
+            data = self.loop.run_until_complete(resp2.json())
+            self.assertEqual(data['headers']['Cookie'],
+                             'c1=cookie1; c2=cookie2')
+
     def test_chunked(self):
         with test_utils.run_server(self.loop, router=Functional) as httpd:
             r = self.loop.run_until_complete(

tests/test_connector.py

@@ -52,6 +52,27 @@ def test_release(self):
         self.connector._release.assert_called_with(
             self.key, self.request, self.transport, self.protocol)
 
+    def test_no_share_cookies(self):
+        connector = aiohttp.BaseConnector(share_cookies=False, loop=self.loop)
+
+        conn = Connection(
+            connector, self.key, self.request,
+            self.transport, self.protocol, self.loop)
+        self.assertEqual(connector.cookies, {})
+        conn.share_cookies({'c1': 'cookie1'})
+        self.assertEqual(connector.cookies, {})
+
+    def test_share_cookies(self):
+        connector = aiohttp.BaseConnector(share_cookies=True, loop=self.loop)
+
+        conn = Connection(
+            connector, self.key, self.request,
+            self.transport, self.protocol, self.loop)
+        self.assertEqual(connector.cookies, {})
+        conn.share_cookies({'c1': 'cookie1'})
+        self.assertEqual(connector.cookies,
+                         http.cookies.SimpleCookie({'c1': 'cookie1'}))
+
 
 class BaseConnectorTests(unittest.TestCase):
 
@@ -123,14 +144,14 @@ def test_release(self, m_time):
         resp = req.response = unittest.mock.Mock()
         resp.message.should_close = False
 
-        cookies = resp.cookies = http.cookies.SimpleCookie()
-        cookies['c1'] = 'cookie1'
-        cookies['c2'] = 'cookie2'
+        # cookies = resp.cookies = http.cookies.SimpleCookie()
+        # cookies['c1'] = 'cookie1'
+        # cookies['c2'] = 'cookie2'
 
         tr, proto = unittest.mock.Mock(), unittest.mock.Mock()
         conn._release(1, req, tr, proto)
         self.assertEqual(conn._conns[1][0], (tr, proto, 10))
-        self.assertEqual(conn.cookies, dict(cookies.items()))
+        # self.assertEqual(conn.cookies, dict(cookies.items()))
         self.assertTrue(conn._start_cleanup_task.called)
 
     def test_release_close(self):
@@ -573,7 +594,8 @@ def test_request_port(self, ClientRequestMock):
         ClientRequestMock.return_value = proxy_req
 
         loop_mock = unittest.mock.Mock()
-        connector = aiohttp.ProxyConnector('http://proxy.example.com', loop=loop_mock)
+        connector = aiohttp.ProxyConnector('http://proxy.example.com',
+                                           loop=loop_mock)
 
         tr, proto = unittest.mock.Mock(), unittest.mock.Mock()
         tr.get_extra_info.return_value = None