From 7922fec9fe2666a171bf9ed97b75ef3c5e7a07d5 Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <abg@uber.com>
Date: Wed, 1 Jul 2015 17:45:28 -0700
Subject: [PATCH] Tornado support

---
 .travis.yml                       |   2 +
 tests/integration/test_tornado.py | 205 ++++++++++++++++++++++++++++++
 tox.ini                           |   5 +-
 vcr/patch.py                      |  62 ++++++++-
 vcr/stubs/tornado_stubs.py        | 147 +++++++++++++++++++++
 5 files changed, 417 insertions(+), 4 deletions(-)
 create mode 100644 tests/integration/test_tornado.py
 create mode 100644 vcr/stubs/tornado_stubs.py

diff --git a/.travis.yml b/.travis.yml
index 93fea8b..15910b1 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -16,6 +16,7 @@ env:
     - WITH_LIB="urllib31.7"
     - WITH_LIB="urllib31.9"
     - WITH_LIB="urllib31.10"
+    - WITH_LIB="tornado"
 matrix:
   allow_failures:
     - env: WITH_LIB="boto"
@@ -45,4 +46,5 @@ install:
 - if [ $WITH_LIB = "urllib31.7" ] ; then pip install certifi urllib3==1.7.1; fi
 - if [ $WITH_LIB = "urllib31.9" ] ; then pip install certifi urllib3==1.9.1; fi
 - if [ $WITH_LIB = "urllib31.10" ] ; then pip install certifi urllib3==1.10.2; fi
+- if [ $WITH_LIB = "tornado" ] ; then pip install tornado==4.2 pytest-tornado pycurl; fi
 script: python setup.py test
diff --git a/tests/integration/test_tornado.py b/tests/integration/test_tornado.py
new file mode 100644
index 0000000..6865924
--- /dev/null
+++ b/tests/integration/test_tornado.py
@@ -0,0 +1,205 @@
+# -*- coding: utf-8 -*-
+'''Test requests' interaction with vcr'''
+
+import json
+
+import pytest
+import vcr
+
+from assertions import assert_cassette_empty, assert_is_json
+
+
+http = pytest.importorskip("tornado.httpclient")
+
+
+@pytest.fixture(params=['simple', 'curl', 'default'])
+def get_client(request):
+    if request.param == 'simple':
+        from tornado import simple_httpclient as simple
+        return (lambda: simple.SimpleAsyncHTTPClient())
+    elif request.param == 'curl':
+        curl = pytest.importorskip("tornado.curl_httpclient")
+        return (lambda: curl.CurlAsyncHTTPClient())
+    else:
+        return (lambda: http.AsyncHTTPClient())
+
+
+def get(client, url, **kwargs):
+    raise_error = kwargs.pop('raise_error', True)
+    return client.fetch(
+        http.HTTPRequest(url, method='GET', **kwargs),
+        raise_error=raise_error,
+    )
+
+
+def post(client, url, data=None, **kwargs):
+    if data:
+        kwargs['body'] = json.dumps(data)
+    return client.fetch(http.HTTPRequest(url, method='POST', **kwargs))
+
+
+@pytest.fixture(params=["https", "http"])
+def scheme(request):
+    '''Fixture that returns both http and https.'''
+    return request.param
+
+
+@pytest.mark.gen_test
+def test_status_code(get_client, scheme, tmpdir):
+    '''Ensure that we can read the status code'''
+    url = scheme + '://httpbin.org/'
+    with vcr.use_cassette(str(tmpdir.join('atts.yaml'))):
+        status_code = (yield get(get_client(), url)).code
+
+    with vcr.use_cassette(str(tmpdir.join('atts.yaml'))) as cass:
+        assert status_code == (yield get(get_client(), url)).code
+        assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_headers(get_client, scheme, tmpdir):
+    '''Ensure that we can read the headers back'''
+    url = scheme + '://httpbin.org/'
+    with vcr.use_cassette(str(tmpdir.join('headers.yaml'))):
+        headers = (yield get(get_client(), url)).headers
+
+    with vcr.use_cassette(str(tmpdir.join('headers.yaml'))) as cass:
+        assert headers == (yield get(get_client(), url)).headers
+        assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_body(get_client, tmpdir, scheme):
+    '''Ensure the responses are all identical enough'''
+
+    url = scheme + '://httpbin.org/bytes/1024'
+    with vcr.use_cassette(str(tmpdir.join('body.yaml'))):
+        content = (yield get(get_client(), url)).body
+
+    with vcr.use_cassette(str(tmpdir.join('body.yaml'))) as cass:
+        assert content == (yield get(get_client(), url)).body
+        assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_auth(get_client, tmpdir, scheme):
+    '''Ensure that we can handle basic auth'''
+    auth = ('user', 'passwd')
+    url = scheme + '://httpbin.org/basic-auth/user/passwd'
+    with vcr.use_cassette(str(tmpdir.join('auth.yaml'))):
+        one = yield get(
+            get_client(), url, auth_username=auth[0], auth_password=auth[1]
+        )
+
+    with vcr.use_cassette(str(tmpdir.join('auth.yaml'))) as cass:
+        two = yield get(
+            get_client(), url, auth_username=auth[0], auth_password=auth[1]
+        )
+        assert one.body == two.body
+        assert one.code == two.code
+        assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_auth_failed(get_client, tmpdir, scheme):
+    '''Ensure that we can save failed auth statuses'''
+    auth = ('user', 'wrongwrongwrong')
+    url = scheme + '://httpbin.org/basic-auth/user/passwd'
+    with vcr.use_cassette(str(tmpdir.join('auth-failed.yaml'))) as cass:
+        # Ensure that this is empty to begin with
+        assert_cassette_empty(cass)
+        one = yield get(
+            get_client(),
+            url,
+            auth_username=auth[0],
+            auth_password=auth[1],
+            raise_error=False
+        )
+
+    with vcr.use_cassette(str(tmpdir.join('auth-failed.yaml'))) as cass:
+        two = yield get(
+            get_client(),
+            url,
+            auth_username=auth[0],
+            auth_password=auth[1],
+            raise_error=False
+        )
+        assert one.body == two.body
+        assert one.code == two.code == 401
+        assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_post(get_client, tmpdir, scheme):
+    '''Ensure that we can post and cache the results'''
+    data = {'key1': 'value1', 'key2': 'value2'}
+    url = scheme + '://httpbin.org/post'
+    with vcr.use_cassette(str(tmpdir.join('requests.yaml'))):
+        req1 = (yield post(get_client(), url, data)).body
+
+    with vcr.use_cassette(str(tmpdir.join('requests.yaml'))) as cass:
+        req2 = (yield post(get_client(), url, data)).body
+
+    assert req1 == req2
+    assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_redirects(get_client, tmpdir, scheme):
+    '''Ensure that we can handle redirects'''
+    url = scheme + '://httpbin.org/redirect-to?url=bytes/1024'
+    with vcr.use_cassette(str(tmpdir.join('requests.yaml'))):
+        content = (yield get(get_client(), url)).body
+
+    with vcr.use_cassette(str(tmpdir.join('requests.yaml'))) as cass:
+        assert content == (yield get(get_client(), url)).body
+        assert cass.play_count == 1
+
+
+@pytest.mark.gen_test
+def test_cross_scheme(get_client, tmpdir, scheme):
+    '''Ensure that requests between schemes are treated separately'''
+    # First fetch a url under http, and then again under https and then
+    # ensure that we haven't served anything out of cache, and we have two
+    # requests / response pairs in the cassette
+    with vcr.use_cassette(str(tmpdir.join('cross_scheme.yaml'))) as cass:
+        yield get(get_client(), 'https://httpbin.org/')
+        yield get(get_client(), 'http://httpbin.org/')
+        assert cass.play_count == 0
+        assert len(cass) == 2
+
+    # Then repeat the same requests and ensure both were replayed.
+    with vcr.use_cassette(str(tmpdir.join('cross_scheme.yaml'))) as cass:
+        yield get(get_client(), 'https://httpbin.org/')
+        yield get(get_client(), 'http://httpbin.org/')
+        assert cass.play_count == 2
+
+
+@pytest.mark.gen_test
+def test_gzip(get_client, tmpdir, scheme):
+    '''
+    Ensure that httpclient is able to automatically decompress the response
+    body
+    '''
+    url = scheme + '://httpbin.org/gzip'
+
+    with vcr.use_cassette(str(tmpdir.join('gzip.yaml'))):
+        response = yield get(get_client(), url, decompress_response=True)
+        assert_is_json(response.body)
+
+    with vcr.use_cassette(str(tmpdir.join('gzip.yaml'))) as cass:
+        response = yield get(get_client(), url, decompress_response=True)
+        assert_is_json(response.body)
+        assert 1 == cass.play_count
+
+
+@pytest.mark.gen_test
+def test_https_with_cert_validation_disabled(get_client, tmpdir):
+    cass_path = str(tmpdir.join('cert_validation_disabled.yaml'))
+
+    with vcr.use_cassette(cass_path):
+        yield get(get_client(), 'https://httpbin.org', validate_cert=False)
+
+    with vcr.use_cassette(cass_path) as cass:
+        yield get(get_client(), 'https://httpbin.org', validate_cert=False)
+        assert 1 == cass.play_count
diff --git a/tox.ini b/tox.ini
index 6230f4b..6f378fa 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = {py26,py27,py33,py34,pypy}-{requests27,requests26,requests25,requests24,requests23,requests22,requests1,httplib2,urllib317,urllib319,urllib3110,boto}
+envlist = {py26,py27,py33,py34,pypy}-{requests27,requests26,requests25,requests24,requests23,requests22,requests1,httplib2,urllib317,urllib319,urllib3110,tornado,boto}
 
 [testenv]
 commands =
@@ -26,4 +26,7 @@ deps =
     urllib317: urllib3==1.7.1
     urllib319: urllib3==1.9.1
     urllib3110: urllib3==1.10.2
+    tornado: tornado
+    tornado: pytest-tornado
+    tornado: pycurl
     boto: boto
diff --git a/vcr/patch.py b/vcr/patch.py
index caed9fa..f34b3fe 100644
--- a/vcr/patch.py
+++ b/vcr/patch.py
@@ -52,6 +52,25 @@ else:
     _CertValidatingHTTPSConnection = boto.https_connection.CertValidatingHTTPSConnection
 
 
+# Try to save the original types for Tornado
+try:
+    import tornado.httpclient
+    import tornado.simple_httpclient
+except ImportError:  # pragma: no cover
+    pass
+else:
+    _AsyncHTTPClient = tornado.httpclient.AsyncHTTPClient
+    _SimpleAsyncHTTPClient = tornado.simple_httpclient.SimpleAsyncHTTPClient
+
+
+try:
+    import tornado.curl_httpclient
+except ImportError:  # pragma: no cover
+    pass
+else:
+    _CurlAsyncHTTPClient = tornado.curl_httpclient.CurlAsyncHTTPClient
+
+
 class CassettePatcherBuilder(object):
 
     def _build_patchers_from_mock_triples_decorator(function):
@@ -68,10 +87,11 @@ class CassettePatcherBuilder(object):
 
     def build(self):
         return itertools.chain(
-            self._httplib(), self._requests(), self._urllib3(), self._httplib2(),
-            self._boto(), self._build_patchers_from_mock_triples(
+            self._httplib(), self._requests(), self._urllib3(),
+            self._httplib2(), self._boto(), self._tornado(),
+            self._build_patchers_from_mock_triples(
                 self._cassette.custom_patches
-            )
+            ),
         )
 
     def _build_patchers_from_mock_triples(self, mock_triples):
@@ -205,6 +225,27 @@ class CassettePatcherBuilder(object):
             from .stubs.boto_stubs import VCRCertValidatingHTTPSConnection
             yield cpool, 'CertValidatingHTTPSConnection', VCRCertValidatingHTTPSConnection
 
+    @_build_patchers_from_mock_triples_decorator
+    def _tornado(self):
+        try:
+            import tornado.httpclient as http
+            import tornado.simple_httpclient as simple
+        except ImportError:  # pragma: no cover
+            pass
+        else:
+            from .stubs.tornado_stubs import VCRAsyncHTTPClient
+            from .stubs.tornado_stubs import VCRSimpleAsyncHTTPClient
+
+            yield http, 'AsyncHTTPClient', VCRAsyncHTTPClient
+            yield simple, 'SimpleAsyncHTTPClient', VCRSimpleAsyncHTTPClient
+        try:
+            import tornado.curl_httpclient as curl
+        except ImportError:  # pragma: no cover
+            pass
+        else:
+            from .stubs.tornado_stubs import VCRCurlAsyncHTTPClient
+            yield curl, 'CurlAsyncHTTPClient', VCRCurlAsyncHTTPClient
+
     def _urllib3_patchers(self, cpool, stubs):
         http_connection_remover = ConnectionRemover(
             self._get_cassette_subclass(stubs.VCRRequestsHTTPConnection)
@@ -320,6 +361,21 @@ def reset_patchers():
         yield mock.patch.object(cpool, 'CertValidatingHTTPSConnection',
                                 _CertValidatingHTTPSConnection)
 
+    try:
+        import tornado.httpclient as http
+        import tornado.simple_httpclient as simple
+    except ImportError:  # pragma: no cover
+        pass
+    else:
+        yield mock.patch.object(http, 'AsyncHTTPClient', _AsyncHTTPClient)
+        yield mock.patch.object(simple, 'SimpleAsyncHTTPClient', _SimpleAsyncHTTPClient)
+    try:
+        import tornado.curl_httpclient as curl
+    except ImportError:  # pragma: no cover
+        pass
+    else:
+        yield mock.patch.object(curl, 'CurlAsyncHTTPClient', _CurlAsyncHTTPClient)
+
 
 @contextlib.contextmanager
 def force_reset():
diff --git a/vcr/stubs/tornado_stubs.py b/vcr/stubs/tornado_stubs.py
new file mode 100644
index 0000000..c6f1bef
--- /dev/null
+++ b/vcr/stubs/tornado_stubs.py
@@ -0,0 +1,147 @@
+'''Stubs for tornado HTTP clients'''
+from __future__ import absolute_import
+
+from six import BytesIO
+
+from tornado import httputil
+from tornado.httpclient import AsyncHTTPClient
+from tornado.httpclient import HTTPResponse
+from tornado.simple_httpclient import SimpleAsyncHTTPClient
+
+from vcr.errors import CannotOverwriteExistingCassetteException
+from vcr.request import Request
+
+
+class _VCRAsyncClient(object):
+    cassette = None
+
+    def __new__(cls, *args, **kwargs):
+        from vcr.patch import force_reset
+        with force_reset():
+            return super(_VCRAsyncClient, cls).__new__(cls, *args, **kwargs)
+
+    def initialize(self, *args, **kwargs):
+        from vcr.patch import force_reset
+        with force_reset():
+            self.real_client = self._baseclass(*args, **kwargs)
+
+    @property
+    def io_loop(self):
+        return self.real_client.io_loop
+
+    @property
+    def _closed(self):
+        return self.real_client._closed
+
+    @property
+    def defaults(self):
+        return self.real_client.defaults
+
+    def close(self):
+        from vcr.patch import force_reset
+        with force_reset():
+            self.real_client.close()
+
+    def fetch_impl(self, request, callback):
+        headers = dict(request.headers)
+        if request.user_agent:
+            headers.setdefault('User-Agent', request.user_agent)
+
+        # TODO body_producer, header_callback, and streaming_callback are not
+        # yet supported.
+
+        unsupported_call = (
+            request.body_producer is not None or
+            request.header_callback is not None or
+            request.streaming_callback is not None
+        )
+        if unsupported_call:
+            response = HTTPResponse(
+                request,
+                599,
+                error=Exception(
+                    "The request (%s) uses AsyncHTTPClient functionality "
+                    "that is not yet supported by VCR.py. Please make the "
+                    "request outside a VCR.py context." % repr(request)
+                ),
+            )
+
+        vcr_request = Request(
+            request.method,
+            request.url,
+            request.body,
+            headers,
+        )
+
+        if self.cassette.can_play_response_for(vcr_request):
+            vcr_response = self.cassette.play_response(vcr_request)
+            headers = httputil.HTTPHeaders()
+
+            recorded_headers = vcr_response['headers']
+            if isinstance(recorded_headers, dict):
+                recorded_headers = recorded_headers.items()
+            for k, vs in recorded_headers:
+                for v in vs:
+                    headers.add(k, v)
+            response = HTTPResponse(
+                request,
+                code=vcr_response['status']['code'],
+                reason=vcr_response['status']['message'],
+                headers=headers,
+                buffer=BytesIO(vcr_response['body']['string']),
+            )
+            callback(response)
+        else:
+            if self.cassette.write_protected and self.cassette.filter_request(
+                vcr_request
+            ):
+                response = HTTPResponse(
+                    request,
+                    599,
+                    error=CannotOverwriteExistingCassetteException(
+                        "No match for the request (%r) was found. "
+                        "Can't overwrite existing cassette (%r) in "
+                        "your current record mode (%r)."
+                        % (vcr_request, self.cassette._path,
+                           self.cassette.record_mode)
+                    ),
+                )
+                callback(response)
+
+            def new_callback(response):
+                headers = [
+                    (k, response.headers.get_list(k))
+                    for k in response.headers.keys()
+                ]
+
+                vcr_response = {
+                    'status': {
+                        'code': response.code,
+                        'message': response.reason,
+                    },
+                    'headers': headers,
+                    'body': {'string': response.body},
+                }
+                self.cassette.append(vcr_request, vcr_response)
+                callback(response)
+
+            from vcr.patch import force_reset
+            with force_reset():
+                self.real_client.fetch_impl(request, new_callback)
+
+
+class VCRAsyncHTTPClient(_VCRAsyncClient, AsyncHTTPClient):
+    _baseclass = AsyncHTTPClient
+
+
+class VCRSimpleAsyncHTTPClient(_VCRAsyncClient, SimpleAsyncHTTPClient):
+    _baseclass = SimpleAsyncHTTPClient
+
+
+try:
+    from tornado.curl_httpclient import CurlAsyncHTTPClient
+except ImportError:  # pragma: no cover
+    VCRCurlAsyncHTTPClient = None
+else:
+    class VCRCurlAsyncHTTPClient(_VCRAsyncClient, CurlAsyncHTTPClient):
+        _baseclass = CurlAsyncHTTPClient