mirror of
https://github.com/kevin1024/vcrpy.git
synced 2025-12-10 09:35:34 +00:00
Use mock for patching http connection objects.
This commit is contained in:
Ivan Malison
@@ -1,4 +1,5 @@
|
||||
'''The container for recorded requests and responses'''
|
||||
import contextlib2
|
||||
import functools
|
||||
try:
|
||||
from collections import Counter
|
||||
@@ -6,7 +7,7 @@ except ImportError:
|
||||
from .compat.counter import Counter
|
||||
|
||||
# Internal imports
|
||||
from .patch import install, reset
|
||||
from .patch import build_patchers
|
||||
from .persist import load_cassette, save_cassette
|
||||
from .filters import filter_request
|
||||
from .serializers import yamlserializer
|
||||
@@ -31,14 +32,27 @@ class CassetteContextDecorator(object):
|
||||
def __init__(self, cls, args_getter):
|
||||
self.cls = cls
|
||||
self._args_getter = args_getter
|
||||
self.__finish = None
|
||||
|
||||
def _patch_generator(self, cassette):
|
||||
with contextlib2.ExitStack() as exit_stack:
|
||||
for patcher in build_patchers(cassette):
|
||||
exit_stack.enter_context(patcher)
|
||||
yield
|
||||
# TODO(@IvanMalison): Hmmm. it kind of feels like this should be somewhere else.
|
||||
cassette._save()
|
||||
|
||||
def __enter__(self):
|
||||
assert self.__finish is None
|
||||
path, kwargs = self._args_getter()
|
||||
self._cassette = self.cls.load(path, **kwargs)
|
||||
return self._cassette.__enter__()
|
||||
cassette = self.cls.load(path, **kwargs)
|
||||
self.__finish = self._patch_generator(cassette)
|
||||
self.__finish.__next__()
|
||||
return cassette
|
||||
|
||||
def __exit__(self, *args):
|
||||
return self._cassette.__exit__(*args)
|
||||
[_ for _ in self.__finish] # this exits the context created by the call to _patch_generator.
|
||||
self.__finish = None
|
||||
|
||||
def __call__(self, function):
|
||||
@functools.wraps(function)
|
||||
@@ -66,17 +80,10 @@ class Cassette(object):
|
||||
def use(cls, *args, **kwargs):
|
||||
return CassetteContextDecorator.from_args(cls, *args, **kwargs)
|
||||
|
||||
def __init__(self,
|
||||
path,
|
||||
serializer=yamlserializer,
|
||||
record_mode='once',
|
||||
match_on=(uri, method),
|
||||
filter_headers=(),
|
||||
filter_query_parameters=(),
|
||||
before_record=None,
|
||||
before_record_response=None,
|
||||
ignore_hosts=(),
|
||||
ignore_localhost=()):
|
||||
def __init__(self, path, serializer=yamlserializer, record_mode='once',
|
||||
match_on=(uri, method), filter_headers=(),
|
||||
filter_query_parameters=(), before_record=None, before_record_response=None,
|
||||
ignore_hosts=(), ignore_localhost=()):
|
||||
self._path = path
|
||||
self._serializer = serializer
|
||||
self._match_on = match_on
|
||||
@@ -228,12 +235,3 @@ class Cassette(object):
|
||||
for response in self._responses(request):
|
||||
return True
|
||||
return False
|
||||
|
||||
def __enter__(self):
|
||||
'''Patch the fetching libraries we know about'''
|
||||
install(self)
|
||||
return self
|
||||
|
||||
def __exit__(self, typ, value, traceback):
|
||||
self._save()
|
||||
reset()
|
||||
|
||||
152
vcr/patch.py
152
vcr/patch.py
@@ -1,4 +1,6 @@
|
||||
'''Utilities for patching in cassettes'''
|
||||
import contextlib2
|
||||
import mock
|
||||
|
||||
from .stubs import VCRHTTPConnection, VCRHTTPSConnection
|
||||
from six.moves import http_client as httplib
|
||||
@@ -8,139 +10,159 @@ from six.moves import http_client as httplib
|
||||
_HTTPConnection = httplib.HTTPConnection
|
||||
_HTTPSConnection = httplib.HTTPSConnection
|
||||
|
||||
|
||||
# Try to save the original types for requests
|
||||
try:
|
||||
# Try to save the original types for requests
|
||||
import requests.packages.urllib3.connectionpool as cpool
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
_VerifiedHTTPSConnection = cpool.VerifiedHTTPSConnection
|
||||
_cpoolHTTPConnection = cpool.HTTPConnection
|
||||
_cpoolHTTPSConnection = cpool.HTTPSConnection
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
# Try to save the original types for urllib3
|
||||
try:
|
||||
# Try to save the original types for urllib3
|
||||
import urllib3
|
||||
_VerifiedHTTPSConnection = urllib3.connectionpool.VerifiedHTTPSConnection
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
_VerifiedHTTPSConnection = urllib3.connectionpool.VerifiedHTTPSConnection
|
||||
|
||||
|
||||
# Try to save the original types for httplib2
|
||||
try:
|
||||
# Try to save the original types for httplib2
|
||||
import httplib2
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
_HTTPConnectionWithTimeout = httplib2.HTTPConnectionWithTimeout
|
||||
_HTTPSConnectionWithTimeout = httplib2.HTTPSConnectionWithTimeout
|
||||
_SCHEME_TO_CONNECTION = httplib2.SCHEME_TO_CONNECTION
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
# Try to save the original types for boto
|
||||
try:
|
||||
# Try to save the original types for boto
|
||||
import boto.https_connection
|
||||
_CertValidatingHTTPSConnection = \
|
||||
boto.https_connection.CertValidatingHTTPSConnection
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
_CertValidatingHTTPSConnection = boto.https_connection.CertValidatingHTTPSConnection
|
||||
|
||||
|
||||
def install(cassette):
|
||||
|
||||
def cassette_subclass(base_class, cassette):
|
||||
return type('{0}{1}'.format(base_class.__name__, cassette._path),
|
||||
(base_class,), dict(cassette=cassette))
|
||||
|
||||
|
||||
def build_patchers(cassette):
|
||||
"""
|
||||
Patch all the HTTPConnections references we can find!
|
||||
Build patches for all the HTTPConnections references we can find!
|
||||
This replaces the actual HTTPConnection with a VCRHTTPConnection
|
||||
object which knows how to save to / read from cassettes
|
||||
"""
|
||||
httplib.HTTPConnection = VCRHTTPConnection
|
||||
httplib.HTTPSConnection = VCRHTTPSConnection
|
||||
httplib.HTTPConnection.cassette = cassette
|
||||
httplib.HTTPSConnection.cassette = cassette
|
||||
_VCRHTTPConnection = cassette_subclass(VCRHTTPConnection, cassette)
|
||||
_VCRHTTPSConnection = cassette_subclass(VCRHTTPSConnection, cassette)
|
||||
|
||||
# patch requests v1.x
|
||||
|
||||
yield mock.patch.object(httplib, 'HTTPConnection', _VCRHTTPConnection)
|
||||
yield mock.patch.object(httplib, 'HTTPSConnection', _VCRHTTPSConnection)
|
||||
|
||||
# requests
|
||||
try:
|
||||
import requests.packages.urllib3.connectionpool as cpool
|
||||
from .stubs.requests_stubs import VCRRequestsHTTPConnection, VCRRequestsHTTPSConnection
|
||||
cpool.VerifiedHTTPSConnection = VCRRequestsHTTPSConnection
|
||||
cpool.HTTPConnection = VCRRequestsHTTPConnection
|
||||
cpool.VerifiedHTTPSConnection.cassette = cassette
|
||||
cpool.HTTPConnection = VCRHTTPConnection
|
||||
cpool.HTTPConnection.cassette = cassette
|
||||
# patch requests v2.x
|
||||
cpool.HTTPConnectionPool.ConnectionCls = VCRRequestsHTTPConnection
|
||||
cpool.HTTPConnectionPool.cassette = cassette
|
||||
cpool.HTTPSConnectionPool.ConnectionCls = VCRRequestsHTTPSConnection
|
||||
cpool.HTTPSConnectionPool.cassette = cassette
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
from .stubs.requests_stubs import VCRRequestsHTTPConnection, VCRRequestsHTTPSConnection
|
||||
# patch requests v1.x
|
||||
yield mock.patch.object(cpool, 'VerifiedHTTPSConnection', cassette_subclass(VCRRequestsHTTPSConnection, cassette))
|
||||
yield mock.patch.object(cpool, 'HTTPConnection', cassette_subclass(VCRRequestsHTTPConnection, cassette))
|
||||
yield mock.patch.object(cpool, 'HTTPConnection', _VCRHTTPConnection)
|
||||
# patch requests v2.x
|
||||
yield mock.patch.object(cpool.HTTPConnectionPool, 'ConnectionCls', cassette_subclass(VCRRequestsHTTPConnection, cassette))
|
||||
yield mock.patch.object(cpool.HTTPSConnectionPool, 'ConnectionCls', cassette_subclass(VCRRequestsHTTPSConnection, cassette))
|
||||
|
||||
# patch urllib3
|
||||
try:
|
||||
import urllib3.connectionpool as cpool
|
||||
from .stubs.urllib3_stubs import VCRVerifiedHTTPSConnection
|
||||
cpool.VerifiedHTTPSConnection = VCRVerifiedHTTPSConnection
|
||||
cpool.VerifiedHTTPSConnection.cassette = cassette
|
||||
cpool.HTTPConnection = VCRHTTPConnection
|
||||
cpool.HTTPConnection.cassette = cassette
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
from .stubs.urllib3_stubs import VCRVerifiedHTTPSConnection
|
||||
yield mock.patch.object(cpool, 'VerifiedHTTPSConnection', cassette_subclass(VCRVerifiedHTTPSConnection, cassette))
|
||||
yield mock.patch.object(cpool, 'HTTPConnection', _VCRHTTPConnection)
|
||||
|
||||
# patch httplib2
|
||||
try:
|
||||
import httplib2 as cpool
|
||||
from .stubs.httplib2_stubs import VCRHTTPConnectionWithTimeout
|
||||
from .stubs.httplib2_stubs import VCRHTTPSConnectionWithTimeout
|
||||
cpool.HTTPConnectionWithTimeout = VCRHTTPConnectionWithTimeout
|
||||
cpool.HTTPSConnectionWithTimeout = VCRHTTPSConnectionWithTimeout
|
||||
cpool.SCHEME_TO_CONNECTION = {
|
||||
'http': VCRHTTPConnectionWithTimeout,
|
||||
'https': VCRHTTPSConnectionWithTimeout
|
||||
}
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
from .stubs.httplib2_stubs import VCRHTTPConnectionWithTimeout
|
||||
from .stubs.httplib2_stubs import VCRHTTPSConnectionWithTimeout
|
||||
yield mock.patch.object(cpool, 'HTTPConnectionWithTimeout', cassette_subclass(VCRHTTPConnectionWithTimeout, cassette))
|
||||
yield mock.patch.object(cpool, 'HTTPSConnectionWithTimeout', cassette_subclass(VCRHTTPSConnectionWithTimeout, cassette))
|
||||
yield mock.patch.object(cpool, 'SCHEME_TO_CONNECTION', {'http': VCRHTTPConnectionWithTimeout, 'https': VCRHTTPSConnectionWithTimeout})
|
||||
|
||||
# patch boto
|
||||
try:
|
||||
import boto.https_connection as cpool
|
||||
from .stubs.boto_stubs import VCRCertValidatingHTTPSConnection
|
||||
cpool.CertValidatingHTTPSConnection = VCRCertValidatingHTTPSConnection
|
||||
cpool.CertValidatingHTTPSConnection.cassette = cassette
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
from .stubs.boto_stubs import VCRCertValidatingHTTPSConnection
|
||||
yield mock.patch.object(cpool, 'CertValidatingHTTPSConnection', cassette_subclass(VCRCertValidatingHTTPSConnection, cassette))
|
||||
|
||||
|
||||
def reset():
|
||||
'''Undo all the patching'''
|
||||
httplib.HTTPConnection = _HTTPConnection
|
||||
httplib.HTTPSConnection = _HTTPSConnection
|
||||
def reset_patchers():
|
||||
yield mock.patch.object(httplib, 'HTTPConnection', _HTTPConnection)
|
||||
yield mock.patch.object(httplib, 'HTTPSConnection', _HTTPSConnection)
|
||||
try:
|
||||
import requests.packages.urllib3.connectionpool as cpool
|
||||
# unpatch requests v1.x
|
||||
cpool.VerifiedHTTPSConnection = _VerifiedHTTPSConnection
|
||||
cpool.HTTPConnection = _cpoolHTTPConnection
|
||||
# unpatch requests v2.x
|
||||
cpool.HTTPConnectionPool.ConnectionCls = _cpoolHTTPConnection
|
||||
cpool.HTTPSConnection = _cpoolHTTPSConnection
|
||||
cpool.HTTPSConnectionPool.ConnectionCls = _cpoolHTTPSConnection
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
# unpatch requests v1.x
|
||||
yield mock.patch.object(cpool, 'VerifiedHTTPSConnection', _VerifiedHTTPSConnection)
|
||||
yield mock.patch.object(cpool, 'HTTPConnection', _cpoolHTTPConnection)
|
||||
# unpatch requests v2.x
|
||||
yield mock.patch.object(cpool.HTTPConnectionPool, 'ConnectionCls', _cpoolHTTPConnection)
|
||||
yield mock.patch.object(cpool, 'HTTPSConnection', _cpoolHTTPSConnection)
|
||||
yield mock.patch.object(cpool.HTTPSConnectionPool, 'ConnectionCls', _cpoolHTTPSConnection)
|
||||
|
||||
try:
|
||||
import urllib3.connectionpool as cpool
|
||||
cpool.VerifiedHTTPSConnection = _VerifiedHTTPSConnection
|
||||
cpool.HTTPConnection = _HTTPConnection
|
||||
cpool.HTTPSConnection = _HTTPSConnection
|
||||
cpool.HTTPConnectionPool.ConnectionCls = _HTTPConnection
|
||||
cpool.HTTPSConnectionPool.ConnectionCls = _HTTPSConnection
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
yield mock.patch.object(cpool, 'VerifiedHTTPSConnection', _VerifiedHTTPSConnection)
|
||||
yield mock.patch.object(cpool, 'HTTPConnection', _HTTPConnection)
|
||||
yield mock.patch.object(cpool, 'HTTPSConnection', _HTTPSConnection)
|
||||
yield mock.patch.object(cpool.HTTPConnectionPool, 'ConnectionCls', _HTTPConnection)
|
||||
yield mock.patch.object(cpool.HTTPSConnectionPool, 'ConnectionCls', _HTTPSConnection)
|
||||
|
||||
try:
|
||||
import httplib2 as cpool
|
||||
cpool.HTTPConnectionWithTimeout = _HTTPConnectionWithTimeout
|
||||
cpool.HTTPSConnectionWithTimeout = _HTTPSConnectionWithTimeout
|
||||
cpool.SCHEME_TO_CONNECTION = _SCHEME_TO_CONNECTION
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
yield mock.patch.object(cpool, 'HTTPConnectionWithTimeout', _HTTPConnectionWithTimeout)
|
||||
yield mock.patch.object(cpool, 'HTTPSConnectionWithTimeout', _HTTPSConnectionWithTimeout)
|
||||
yield mock.patch.object(cpool, 'SCHEME_TO_CONNECTION', _SCHEME_TO_CONNECTION)
|
||||
|
||||
try:
|
||||
import boto.https_connection as cpool
|
||||
cpool.CertValidatingHTTPSConnection = _CertValidatingHTTPSConnection
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
else:
|
||||
yield mock.patch.object(cpool, 'CertValidatingHTTPSConnection', _CertValidatingHTTPSConnection)
|
||||
|
||||
@contextlib2.contextmanager
|
||||
def force_reset():
|
||||
with contextlib2.ExitStack() as exit_stack:
|
||||
for patcher in reset_patchers():
|
||||
exit_stack.enter_context(patcher)
|
||||
yield
|
||||
|
||||
@@ -295,10 +295,9 @@ class VCRConnection:
|
||||
# need to temporarily reset here because the real connection
|
||||
# inherits from the thing that we are mocking out. Take out
|
||||
# the reset if you want to see what I mean :)
|
||||
from vcr.patch import install, reset
|
||||
reset()
|
||||
self.real_connection = self._baseclass(*args, **kwargs)
|
||||
install(self.cassette)
|
||||
from vcr.patch import force_reset
|
||||
with force_reset():
|
||||
self.real_connection = self._baseclass(*args, **kwargs)
|
||||
|
||||
|
||||
class VCRHTTPConnection(VCRConnection):
|
||||
|
||||
Reference in New Issue
Block a user