Skip to content

Commit

Permalink
Fix for #957: allow to override default cookie jar (#963)
Browse files Browse the repository at this point in the history
  • Loading branch information
asvetlov authored Jul 14, 2016
1 parent 6f9d172 commit 47028fa
Show file tree
Hide file tree
Showing 7 changed files with 102 additions and 54 deletions.
21 changes: 21 additions & 0 deletions aiohttp/abc.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import asyncio
import sys
from abc import ABC, abstractmethod
from http.cookies import SimpleCookie


PY_35 = sys.version_info >= (3, 5)
Expand Down Expand Up @@ -68,3 +69,23 @@ def resolve(self, hostname):
@abstractmethod
def close(self):
"""Release resolver"""


class AbstractCookieJar(ABC):

def __init__(self, *, loop=None):
self._cookies = SimpleCookie()
self._loop = loop or asyncio.get_event_loop()

@property
def cookies(self):
"""The session cookies."""
return self._cookies

@abstractmethod
def update_cookies(self, cookies, response_url=None):
"""Update cookies."""

@abstractmethod
def filter_cookies(self, request_url):
"""Returns this jar's cookies filtered by their attributes."""
7 changes: 5 additions & 2 deletions aiohttp/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,8 @@ def __init__(self, *, connector=None, loop=None, cookies=None,
auth=None, request_class=ClientRequest,
response_class=ClientResponse,
ws_response_class=ClientWebSocketResponse,
version=aiohttp.HttpVersion11):
version=aiohttp.HttpVersion11,
cookie_jar=None):

if connector is None:
connector = aiohttp.TCPConnector(loop=loop)
Expand All @@ -52,7 +53,9 @@ def __init__(self, *, connector=None, loop=None, cookies=None,
if loop.get_debug():
self._source_traceback = traceback.extract_stack(sys._getframe(1))

self._cookie_jar = CookieJar(loop=loop)
if cookie_jar is None:
cookie_jar = CookieJar(loop=loop)
self._cookie_jar = cookie_jar

if cookies is not None:
self._cookie_jar.update_cookies(cookies)
Expand Down
16 changes: 4 additions & 12 deletions aiohttp/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import multidict

from . import hdrs
from .abc import AbstractCookieJar
from .errors import InvalidURL
try:
from asyncio import ensure_future
Expand Down Expand Up @@ -568,7 +569,7 @@ def _cancel_task(self):
self._cancelled = self._task.cancel()


class CookieJar:
class CookieJar(AbstractCookieJar):
"""Implements cookie storage adhering to RFC 6265."""

DATE_TOKENS_RE = re.compile(
Expand All @@ -584,19 +585,10 @@ class CookieJar:

DATE_YEAR_RE = re.compile("(\d{2,4})")

def __init__(self, cookies=None, loop=None):
self._cookies = SimpleCookie()
self._loop = loop or asyncio.get_event_loop()
def __init__(self, *, loop=None):
super().__init__(loop=loop)
self._host_only_cookies = set()

if cookies is not None:
self.update_cookies(cookies)

@property
def cookies(self):
"""The session cookies."""
return self._cookies

def _expire_cookie(self, name):
if name in self._cookies:
del self._cookies[name]
Expand Down
19 changes: 18 additions & 1 deletion docs/client_reference.rst
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,9 @@ The client session supports the context manager protocol for self closing.
headers=None, skip_auto_headers=None, \
auth=None, request_class=ClientRequest, \
response_class=ClientResponse, \
ws_response_class=ClientWebSocketResponse)
ws_response_class=ClientWebSocketResponse,
version=aiohttp.HttpVersion11,
cookie_jar=None)

The class for creating client sessions and making requests.

Expand Down Expand Up @@ -93,6 +95,21 @@ The client session supports the context manager protocol for self closing.

.. versionadded:: 0.16

:param version: supported HTTP version, ``HTTP 1.1`` by default.

.. versionadded:: 0.21

:param cookie_jar: Cookie Jar, :class:`AbstractCookieJar` instance.

By default every session instance has own private cookie jar for
automatic cookies processing but user may redefine this behavior
by providing own jar implementation.

One example is not processing cookies at all when working in
proxy mode.

.. versionadded:: 0.22

.. versionchanged:: 0.16
*request_class* default changed from ``None`` to ``ClientRequest``

Expand Down
27 changes: 27 additions & 0 deletions docs/web_abc.rst
Original file line number Diff line number Diff line change
Expand Up @@ -97,3 +97,30 @@ attribute.
.. attribute:: request

:class:`aiohttp.web.Request` instance for performing the request.


Abstract Cookie Jar
-------------------

.. class:: AbstractCookieJar(*, loop=None)

An abstract class for cookie storage.

:param loop: an :ref:`event loop<asyncio-event-loop>` instance.

If param is ``None`` :func:`asyncio.get_event_loop`
used for getting default event loop, but we strongly
recommend to use explicit loops everywhere.


.. attribute:: cookies

:class:`http.cookies.SimpleCookie` instance for storing cookies info.

.. method:: update_cookies(cookies, response_url=None)

Update cookies.

.. method:: filter_cookies(request_url)

Returns this jar's cookies filtered by their attributes.
61 changes: 24 additions & 37 deletions tests/test_client_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -377,16 +377,8 @@ def test_request_ctx_manager_props(loop):
def test_cookie_jar_usage(create_app_and_client):
req_url = None

init_mock = mock.Mock(return_value=None)
update_mock = mock.Mock(return_value=None)
filter_mock = mock.Mock(return_value=None)

patches = mock.patch.multiple(
"aiohttp.helpers.CookieJar",
__init__=init_mock,
update_cookies=update_mock,
filter_cookies=filter_mock,
)
jar = mock.Mock()
jar.filter_cookies.return_value = None

@asyncio.coroutine
def handler(request):
Expand All @@ -397,30 +389,25 @@ def handler(request):
resp.set_cookie("response", "resp_value")
return resp

with patches:
app, client = yield from create_app_and_client(
client_params={"cookies": {"request": "req_value"}}
)
app.router.add_route('GET', '/', handler)

# Updating the cookie jar with initial user defined cookies
assert init_mock.called
assert update_mock.called
assert update_mock.call_args[0] == (
{"request": "req_value"},
)

update_mock.reset_mock()
yield from client.get("/")

# Filtering the cookie jar before sending the request,
# getting the request URL as only parameter
assert filter_mock.called
assert filter_mock.call_args[0] == (req_url,)

# Updating the cookie jar with the response cookies
assert update_mock.called
resp_cookies = update_mock.call_args[0][0]
assert isinstance(resp_cookies, http.cookies.SimpleCookie)
assert "response" in resp_cookies
assert resp_cookies["response"].value == "resp_value"
app, client = yield from create_app_and_client(
client_params={"cookies": {"request": "req_value"},
"cookie_jar": jar}
)
app.router.add_route('GET', '/', handler)

# Updating the cookie jar with initial user defined cookies
jar.update_cookies.assert_called_with({"request": "req_value"})

jar.update_cookies.reset_mock()
yield from client.get("/")

# Filtering the cookie jar before sending the request,
# getting the request URL as only parameter
jar.filter_cookies.assert_called_with(req_url)

# Updating the cookie jar with the response cookies
assert jar.update_cookies.called
resp_cookies = jar.update_cookies.call_args[0][0]
assert isinstance(resp_cookies, http.cookies.SimpleCookie)
assert "response" in resp_cookies
assert resp_cookies["response"].value == "resp_value"
5 changes: 3 additions & 2 deletions tests/test_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -351,9 +351,10 @@ def timed_request(
return cookies_sent

def test_constructor(self):
jar = helpers.CookieJar(self.cookies_to_send, self.loop)
jar = helpers.CookieJar(loop=self.loop)
jar.update_cookies(self.cookies_to_send)
self.assertEqual(jar.cookies, self.cookies_to_send)
self.assertEqual(jar._loop, self.loop)
self.assertIs(jar._loop, self.loop)

def test_domain_filter_ip(self):
cookies_sent, cookies_received = (
Expand Down

0 comments on commit 47028fa

Please sign in to comment.