Merge pull request #266 from lamenezes/aiohttp-support

Aiohttp support

.travis.yml

@@ -22,11 +22,14 @@ env:
     - TOX_SUFFIX="urllib3110"
     - TOX_SUFFIX="tornado3"
     - TOX_SUFFIX="tornado4"
+    - TOX_SUFFIX="aiohttp"
 matrix:
   allow_failures:
     - env: TOX_SUFFIX="boto"
     - env: TOX_SUFFIX="boto3"
   exclude:
+    - env: TOX_SUFFIX="flakes"
+      python: 2.6
     - env: TOX_SUFFIX="boto"
       python: 3.3
     - env: TOX_SUFFIX="boto"
@@ -35,6 +38,16 @@ matrix:
       python: 3.4
     - env: TOX_SUFFIX="requests1"
       python: 3.5
+    - env: TOX_SUFFIX="aiohttp"
+      python: 2.6
+    - env: TOX_SUFFIX="aiohttp"
+      python: 2.7
+    - env: TOX_SUFFIX="aiohttp"
+      python: 3.3
+    - env: TOX_SUFFIX="aiohttp"
+      python: pypy
+    - env: TOX_SUFFIX="aiohttp"
+      python: pypy3
 python:
 - 2.6
 - 2.7

tests/integration/aiohttp_utils.py

@@ -0,0 +1,7 @@
+import asyncio
+
+
+@asyncio.coroutine
+def aiohttp_request(session, method, url, as_text, **kwargs):
+    response = yield from session.request(method, url, **kwargs)  # NOQA: E999
+    return response, (yield from response.text()) if as_text else (yield from response.json())  # NOQA: E999

tests/integration/test_aiohttp.py

@@ -0,0 +1,87 @@
+import pytest
+aiohttp = pytest.importorskip("aiohttp")
+
+import asyncio  # NOQA
+import sys  # NOQA
+
+import aiohttp  # NOQA
+import pytest  # NOQA
+import vcr  # NOQA
+
+from .aiohttp_utils import aiohttp_request  # NOQA
+
+
+def get(url, as_text=True, **kwargs):
+    loop = asyncio.get_event_loop()
+    with aiohttp.ClientSession() as session:
+        task = loop.create_task(aiohttp_request(session, 'GET', url, as_text, **kwargs))
+        return loop.run_until_complete(task)
+
+
+def post(url, as_text=True, **kwargs):
+    loop = asyncio.get_event_loop()
+    with aiohttp.ClientSession() as session:
+        task = loop.create_task(aiohttp_request(session, 'POST', url, as_text, **kwargs))
+        return loop.run_until_complete(task)
+
+
+@pytest.fixture(params=["https", "http"])
+def scheme(request):
+    '''Fixture that returns both http and https.'''
+    return request.param
+
+
+def test_status(tmpdir, scheme):
+    url = scheme + '://httpbin.org'
+    with vcr.use_cassette(str(tmpdir.join('status.yaml'))):
+        response, _ = get(url)
+
+    with vcr.use_cassette(str(tmpdir.join('status.yaml'))) as cassette:
+        cassette_response, _ = get(url)
+        assert cassette_response.status == response.status
+        assert cassette.play_count == 1
+
+
+def test_headers(tmpdir, scheme):
+    url = scheme + '://httpbin.org'
+    with vcr.use_cassette(str(tmpdir.join('headers.yaml'))):
+        response, _ = get(url)
+
+    with vcr.use_cassette(str(tmpdir.join('headers.yaml'))) as cassette:
+        cassette_response, _ = get(url)
+        assert cassette_response.headers == response.headers
+        assert cassette.play_count == 1
+
+
+def test_text(tmpdir, scheme):
+    url = scheme + '://httpbin.org'
+    with vcr.use_cassette(str(tmpdir.join('text.yaml'))):
+        _, response_text = get(url)
+
+    with vcr.use_cassette(str(tmpdir.join('text.yaml'))) as cassette:
+        _, cassette_response_text = get(url)
+        assert cassette_response_text == response_text
+        assert cassette.play_count == 1
+
+
+def test_json(tmpdir, scheme):
+    url = scheme + '://httpbin.org/get'
+    with vcr.use_cassette(str(tmpdir.join('json.yaml'))):
+        _, response_json = get(url, as_text=False)
+
+    with vcr.use_cassette(str(tmpdir.join('json.yaml'))) as cassette:
+        _, cassette_response_json = get(url, as_text=False)
+        assert cassette_response_json == response_json
+        assert cassette.play_count == 1
+
+
+def test_post(tmpdir, scheme):
+    data = {'key1': 'value1', 'key2': 'value2'}
+    url = scheme + '://httpbin.org/post'
+    with vcr.use_cassette(str(tmpdir.join('post.yaml'))):
+        _, response_json = post(url, data=data)
+
+    with vcr.use_cassette(str(tmpdir.join('post.yaml'))) as cassette:
+        _, cassette_response_json = post(url, data=data)
+        assert cassette_response_json == response_json
+        assert cassette.play_count == 1

tox.ini

@@ -1,11 +1,11 @@
 [tox]
-envlist = {py26,py27,py33,py34,pypy,pypy3}-{flakes,requests27,requests26,requests25,requests24,requests23,requests22,requests1,httplib2,urllib317,urllib319,urllib3110,tornado3,tornado4,boto,boto3}
+envlist = {py26,py27,py33,py34,pypy,pypy3}-{flakes,requests27,requests26,requests25,requests24,requests23,requests22,requests1,httplib2,urllib317,urllib319,urllib3110,tornado3,tornado4,boto,boto3,aiohttp}
 
 [testenv:flakes]
 skipsdist = True
 commands =
     flake8 --version
-    flake8 --exclude="./docs/conf.py"
+    flake8 --exclude=./docs/conf.py,./.tox/
     pyflakes ./docs/conf.py
 deps = flake8
 
@@ -38,6 +38,7 @@ deps =
     {py26,py27,py33,py34}-tornado4: pycurl
     boto: boto
     boto3: boto3
+    aiohttp: aiohttp
 
 [flake8]
 max_line_length = 110

vcr/patch.py

@@ -80,6 +80,13 @@
     _CurlAsyncHTTPClient_fetch_impl = \
         tornado.curl_httpclient.CurlAsyncHTTPClient.fetch_impl
 
+try:
+    import aiohttp.client
+except ImportError:  # pragma: no cover
+    pass
+else:
+    _AiohttpClientSessionRequest = aiohttp.client.ClientSession._request
+
 
 class CassettePatcherBuilder(object):
 
@@ -98,7 +105,7 @@ def __init__(self, cassette):
     def build(self):
         return itertools.chain(
             self._httplib(), self._requests(), self._boto3(), self._urllib3(),
-            self._httplib2(), self._boto(), self._tornado(),
+            self._httplib2(), self._boto(), self._tornado(), self._aiohttp(),
             self._build_patchers_from_mock_triples(
                 self._cassette.custom_patches
             ),
@@ -273,6 +280,19 @@ def _tornado(self):
             )
             yield curl.CurlAsyncHTTPClient, 'fetch_impl', new_fetch_impl
 
+    @_build_patchers_from_mock_triples_decorator
+    def _aiohttp(self):
+        try:
+            import aiohttp.client as client
+        except ImportError:  # pragma: no cover
+            pass
+        else:
+            from .stubs.aiohttp_stubs import vcr_request
+            new_request = vcr_request(
+                self._cassette, _AiohttpClientSessionRequest
+            )
+            yield client.ClientSession, '_request', new_request
+
     def _urllib3_patchers(self, cpool, stubs):
         http_connection_remover = ConnectionRemover(
             self._get_cassette_subclass(stubs.VCRRequestsHTTPConnection)

vcr/stubs/aiohttp_stubs.py

@@ -0,0 +1,76 @@
+'''Stubs for aiohttp HTTP clients'''
+from __future__ import absolute_import
+
+import asyncio
+import functools
+import json
+
+from aiohttp import ClientResponse
+
+from vcr.request import Request
+
+
+class MockClientResponse(ClientResponse):
+    # TODO: get encoding from header
+    @asyncio.coroutine
+    def json(self, *, encoding='utf-8', loads=json.loads):  # NOQA: E999
+        return loads(self.content.decode(encoding))
+
+    @asyncio.coroutine
+    def text(self, encoding='utf-8'):
+        return self.content.decode(encoding)
+
+    @asyncio.coroutine
+    def release(self):
+        pass
+
+
+def vcr_request(cassette, real_request):
+
+    @functools.wraps(real_request)
+    @asyncio.coroutine
+    def new_request(self, method, url, **kwargs):
+        headers = kwargs.get('headers')
+        headers = self._prepare_headers(headers)
+        data = kwargs.get('data')
+
+        vcr_request = Request(method, url, data, headers)
+
+        if cassette.can_play_response_for(vcr_request):
+            vcr_response = cassette.play_response(vcr_request)
+
+            response = MockClientResponse(method, vcr_response.get('url'))
+            response.status = vcr_response['status']['code']
+            response.content = vcr_response['body']['string']
+            response.reason = vcr_response['status']['message']
+            response.headers = vcr_response['headers']
+
+            response.close()
+            return response
+
+        if cassette.write_protected and cassette.filter_request(vcr_request):
+            response = MockClientResponse(method, url)
+            response.status = 599
+            msg = ("No match for the request {!r} was found. Can't overwrite "
+                   "existing cassette {!r} in your current record mode {!r}.")
+            msg = msg.format(vcr_request, cassette._path, cassette.record_mode)
+            response.content = msg.encode()
+            response.close()
+            return response
+
+        response = yield from real_request(self, method, url, **kwargs)  # NOQA: E999
+
+        vcr_response = {
+            'status': {
+                'code': response.status,
+                'message': response.reason,
+            },
+            'headers': dict(response.headers),
+            'body': {'string': (yield from response.text())},  # NOQA: E999
+            'url': response.url,
+        }
+        cassette.append(vcr_request, vcr_response)
+
+        return response
+
+    return new_request