diff --git a/docs/response.rst b/docs/response.rst index 5f4ed83..af8fb02 100644 --- a/docs/response.rst +++ b/docs/response.rst @@ -34,6 +34,7 @@ Responses are registered with the :py:meth:`requests_mock.Adapter.register_uri` :status_code: The HTTP status response to return. Defaults to 200. :reason: The reason text that accompanies the Status (e.g. 'OK' in '200 OK') :headers: A dictionary of headers to be included in the response. +:cookies: A CookieJar containing all the cookies to add to the response. To specify the body of the response there are a number of options that depend on the format that you wish to return. @@ -82,6 +83,7 @@ The available properties on the `context` are: :headers: The dictionary of headers that are to be returned in the response. :status_code: The status code that is to be returned in the response. :reason: The string HTTP status code reason that is to be returned in the response. +:cookies: A :py:class:`requests_mock.CookieJar` of cookies that will be merged into the response. These parameters are populated initially from the variables provided to the :py:meth:`~requests_mock.Adapter.register_uri` function and if they are modified on the context object then those changes will be reflected in the response. @@ -130,3 +132,33 @@ Callbacks work within response lists in exactly the same way they do normally; >>> resp = session.get('mock://test.com/5') >>> resp.status_code, resp.headers, resp.text (200, {'Test1': 'value1', 'Test2': 'value2'}, 'response') + +Handling Cookies +================ + +Whilst cookies are just headers they are treated in a different way, both in HTTP and the requests library. +To work as closely to the requests library as possible there are two ways to provide cookies to requests_mock responses. + +The most simple method is to use a dictionary interface. +The Key and value of the dictionary are turned directly into the name and value of the cookie. +This method does not allow you to set any of the more advanced cookie parameters like expiry or domain. + +.. doctest:: + + >>> adapter.register_uri('GET', 'mock://test.com/6', cookies={'foo': 'bar'}), + >>> resp = session.get('mock://test.com/6') + >>> resp.cookies['foo'] + 'bar' + +The more advanced way is to construct and populate a cookie jar that you can add cookies to and pass that to the mocker. + +.. doctest:: + + >>> jar = requests_mock.CookieJar() + >>> jar.set('foo', 'bar', domain='.test.com', path='/baz') + >>> adapter.register_uri('GET', 'mock://test.com/7', cookies=jar), + >>> resp = session.get('mock://test.com/7') + >>> resp.cookies['foo'] + 'bar' + >>> resp.cookies.list_paths() + ['/baz'] diff --git a/requests_mock/__init__.py b/requests_mock/__init__.py index b7d0170..3387ded 100755 --- a/requests_mock/__init__.py +++ b/requests_mock/__init__.py @@ -14,12 +14,13 @@ from requests_mock.adapter import Adapter, ANY from requests_mock.exceptions import MockException, NoMockAddress from requests_mock.mocker import mock, Mocker, MockerCore from requests_mock.mocker import DELETE, GET, HEAD, OPTIONS, PATCH, POST, PUT -from requests_mock.response import create_response +from requests_mock.response import create_response, CookieJar __all__ = ['Adapter', 'ANY', 'create_response', + 'CookieJar', 'mock', 'Mocker', 'MockerCore', diff --git a/requests_mock/response.py b/requests_mock/response.py index d04d66d..54841fa 100644 --- a/requests_mock/response.py +++ b/requests_mock/response.py @@ -13,6 +13,9 @@ import json as jsonutils from requests.adapters import HTTPAdapter +from requests.cookies import MockRequest, MockResponse +from requests.cookies import RequestsCookieJar +from requests.cookies import merge_cookies, cookiejar_from_dict from requests.packages.urllib3.response import HTTPResponse import six @@ -20,12 +23,40 @@ from requests_mock import compat from requests_mock import exceptions _BODY_ARGS = frozenset(['raw', 'body', 'content', 'text', 'json']) -_HTTP_ARGS = frozenset(['status_code', 'reason', 'headers']) +_HTTP_ARGS = frozenset(['status_code', 'reason', 'headers', 'cookies']) _DEFAULT_STATUS = 200 _http_adapter = HTTPAdapter() +class CookieJar(RequestsCookieJar): + + def set(self, name, value, **kwargs): + """Add a cookie to the Jar. + + :param str name: cookie name/key. + :param str value: cookie value. + :param int version: Integer or None. Netscape cookies have version 0. + RFC 2965 and RFC 2109 cookies have a version cookie-attribute of 1. + However, note that cookielib may 'downgrade' RFC 2109 cookies to + Netscape cookies, in which case version is 0. + :param str port: String representing a port or a set of ports + (eg. '80', or '80,8080'), + :param str domain: The domain the cookie should apply to. + :param str path: Cookie path (a string, eg. '/acme/rocket_launchers'). + :param bool secure: True if cookie should only be returned over a + secure connection. + :param int expires: Integer expiry date in seconds since epoch or None. + :param bool discard: True if this is a session cookie. + :param str comment: String comment from the server explaining the + function of this cookie. + :param str comment_url: URL linking to a comment from the server + explaining the function of this cookie. + """ + # just here to provide the function documentation + return super(CookieJar, self).set(name, value, **kwargs) + + def _check_body_arguments(**kwargs): # mutual exclusion, only 1 body method may be provided provided = [x for x in _BODY_ARGS if kwargs.pop(x, None) is not None] @@ -53,6 +84,25 @@ class _FakeConnection(object): pass +def _extract_cookies(request, response, cookies): + """Add cookies to the response. + + Cookies in requests are extracted from the headers in the original_response + httplib.HTTPMessage which we don't create so we have to do this step + manually. + """ + # This will add cookies set manually via the Set-Cookie or Set-Cookie2 + # header but this only allows 1 cookie to be set. + http_message = compat._FakeHTTPMessage(response.headers) + response.cookies.extract_cookies(MockResponse(http_message), + MockRequest(request)) + + # This allows you to pass either a CookieJar or a dictionary to request_uri + # or directly to create_response. To allow more than one cookie to be set. + if cookies: + merge_cookies(response.cookies, cookies) + + def create_response(request, **kwargs): """ :param int status_code: The status code to return upon a successful @@ -67,6 +117,8 @@ def create_response(request, **kwargs): and returned upon a successful match. :param dict headers: A dictionary object containing headers that are returned upon a successful match. + :param CookieJar cookies: A cookie jar with cookies to set on the + response. """ connection = kwargs.pop('connection', _FakeConnection()) @@ -103,16 +155,20 @@ def create_response(request, **kwargs): response = _http_adapter.build_response(request, raw) response.connection = connection response.encoding = encoding + + _extract_cookies(request, response, kwargs.get('cookies')) + return response class _Context(object): """Stores the data being used to process a current URL match.""" - def __init__(self, headers, status_code, reason): + def __init__(self, headers, status_code, reason, cookies): self.headers = headers self.status_code = status_code self.reason = reason + self.cookies = cookies class _MatcherResponse(object): @@ -148,9 +204,16 @@ class _MatcherResponse(object): if self._exc: raise self._exc + # If a cookie dict is passed convert it into a CookieJar so that the + # cookies object available in a callback context is always a jar. + cookies = self._params.get('cookies', CookieJar()) + if isinstance(cookies, dict): + cookies = cookiejar_from_dict(cookies, CookieJar()) + context = _Context(self._params.get('headers', {}).copy(), self._params.get('status_code', _DEFAULT_STATUS), - self._params.get('reason')) + self._params.get('reason'), + cookies) # if a body element is a callback then execute it def _call(f, *args, **kwargs): @@ -164,4 +227,5 @@ class _MatcherResponse(object): raw=self._params.get('raw'), status_code=context.status_code, reason=context.reason, - headers=context.headers) + headers=context.headers, + cookies=context.cookies) diff --git a/requests_mock/tests/test_adapter.py b/requests_mock/tests/test_adapter.py index 0231cbc..7c59ac5 100644 --- a/requests_mock/tests/test_adapter.py +++ b/requests_mock/tests/test_adapter.py @@ -488,3 +488,91 @@ class SessionAdapterTests(base.TestCase): self.assertEqual(self.url, self.adapter.last_request.url) self.assertIs(m, self.adapter.last_request.matcher) + + def test_cookies_from_header(self): + headers = {'Set-Cookie': 'fig=newton; Path=/test; domain=.example.com'} + self.adapter.register_uri('GET', + self.url, + text='text', + headers=headers) + + resp = self.session.get(self.url) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual(['/test'], resp.cookies.list_paths()) + self.assertEqual(['.example.com'], resp.cookies.list_domains()) + + def test_cookies_from_dict(self): + # This is a syntax we get from requests. I'm not sure i like it. + self.adapter.register_uri('GET', + self.url, + text='text', + cookies={'fig': 'newton', 'sugar': 'apple'}) + + resp = self.session.get(self.url) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual('apple', resp.cookies['sugar']) + + def test_cookies_with_jar(self): + jar = requests_mock.CookieJar() + jar.set('fig', 'newton', path='/foo', domain='.example.com') + jar.set('sugar', 'apple', path='/bar', domain='.example.com') + + self.adapter.register_uri('GET', self.url, text='text', cookies=jar) + resp = self.session.get(self.url) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual('apple', resp.cookies['sugar']) + self.assertEqual(set(['/foo', '/bar']), set(resp.cookies.list_paths())) + self.assertEqual(['.example.com'], resp.cookies.list_domains()) + + def test_cookies_header_with_cb(self): + + def _cb(request, context): + val = 'fig=newton; Path=/test; domain=.example.com' + context.headers['Set-Cookie'] = val + return 'text' + + self.adapter.register_uri('GET', self.url, text=_cb) + resp = self.session.get(self.url) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual(['/test'], resp.cookies.list_paths()) + self.assertEqual(['.example.com'], resp.cookies.list_domains()) + + def test_cookies_from_dict_with_cb(self): + def _cb(request, context): + # converted into a jar by now + context.cookies.set('sugar', 'apple', path='/test') + return 'text' + + self.adapter.register_uri('GET', + self.url, + text=_cb, + cookies={'fig': 'newton'}) + + resp = self.session.get(self.url) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual('apple', resp.cookies['sugar']) + self.assertEqual(['/', '/test'], resp.cookies.list_paths()) + + def test_cookies_with_jar_cb(self): + def _cb(request, context): + context.cookies.set('sugar', + 'apple', + path='/bar', + domain='.example.com') + return 'text' + + jar = requests_mock.CookieJar() + jar.set('fig', 'newton', path='/foo', domain='.example.com') + + self.adapter.register_uri('GET', self.url, text=_cb, cookies=jar) + resp = self.session.get(self.url) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual('apple', resp.cookies['sugar']) + self.assertEqual(set(['/foo', '/bar']), set(resp.cookies.list_paths())) + self.assertEqual(['.example.com'], resp.cookies.list_domains()) diff --git a/requests_mock/tests/test_response.py b/requests_mock/tests/test_response.py index a571e7f..42e00a5 100644 --- a/requests_mock/tests/test_response.py +++ b/requests_mock/tests/test_response.py @@ -77,3 +77,31 @@ class ResponseTests(base.TestCase): resp = self.create_response() self.assertRaises(exceptions.InvalidRequest, resp.connection.send, self.request) + + def test_cookies_from_header(self): + # domain must be same as request url to pass policy check + headers = {'Set-Cookie': 'fig=newton; Path=/test; domain=.test.url'} + resp = self.create_response(headers=headers) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual(['/test'], resp.cookies.list_paths()) + self.assertEqual(['.test.url'], resp.cookies.list_domains()) + + def test_cookies_from_dict(self): + # This is a syntax we get from requests. I'm not sure i like it. + resp = self.create_response(cookies={'fig': 'newton', + 'sugar': 'apple'}) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual('apple', resp.cookies['sugar']) + + def test_cookies_with_jar(self): + jar = response.CookieJar() + jar.set('fig', 'newton', path='/foo', domain='.test.url') + jar.set('sugar', 'apple', path='/bar', domain='.test.url') + resp = self.create_response(cookies=jar) + + self.assertEqual('newton', resp.cookies['fig']) + self.assertEqual('apple', resp.cookies['sugar']) + self.assertEqual(set(['/foo', '/bar']), set(resp.cookies.list_paths())) + self.assertEqual(['.test.url'], resp.cookies.list_domains())