mirror of
https://github.com/kevin1024/vcrpy.git
synced 2025-12-09 01:03:24 +00:00
Revert "Remove ConnectionRemover class that tried to get rid of vcr connections in ConnectionPools."
This reverts commit dc249b0965.
Conflicts:
vcr/patch.py
This commit is contained in:
61
vcr/patch.py
61
vcr/patch.py
@@ -123,6 +123,12 @@ class CassettePatcherBuilder(object):
|
|||||||
except ImportError: # pragma: no cover
|
except ImportError: # pragma: no cover
|
||||||
return ()
|
return ()
|
||||||
from .stubs.requests_stubs import VCRRequestsHTTPConnection, VCRRequestsHTTPSConnection
|
from .stubs.requests_stubs import VCRRequestsHTTPConnection, VCRRequestsHTTPSConnection
|
||||||
|
http_connection_remover = ConnectionRemover(
|
||||||
|
self._get_cassette_subclass(VCRHTTPConnection)
|
||||||
|
)
|
||||||
|
https_connection_remover = ConnectionRemover(
|
||||||
|
self._get_cassette_subclass(VCRHTTPSConnection)
|
||||||
|
)
|
||||||
mock_triples = (
|
mock_triples = (
|
||||||
(cpool, 'VerifiedHTTPSConnection', VCRRequestsHTTPSConnection),
|
(cpool, 'VerifiedHTTPSConnection', VCRRequestsHTTPSConnection),
|
||||||
(cpool, 'VerifiedHTTPSConnection', VCRRequestsHTTPSConnection),
|
(cpool, 'VerifiedHTTPSConnection', VCRRequestsHTTPSConnection),
|
||||||
@@ -130,16 +136,24 @@ class CassettePatcherBuilder(object):
|
|||||||
(cpool, 'HTTPSConnection', VCRRequestsHTTPSConnection),
|
(cpool, 'HTTPSConnection', VCRRequestsHTTPSConnection),
|
||||||
(cpool.HTTPConnectionPool, 'ConnectionCls', VCRRequestsHTTPConnection),
|
(cpool.HTTPConnectionPool, 'ConnectionCls', VCRRequestsHTTPConnection),
|
||||||
(cpool.HTTPSConnectionPool, 'ConnectionCls', VCRRequestsHTTPSConnection),
|
(cpool.HTTPSConnectionPool, 'ConnectionCls', VCRRequestsHTTPSConnection),
|
||||||
# These handle making sure that sessions only use the
|
|
||||||
# connections of the appropriate type.
|
|
||||||
)
|
)
|
||||||
|
# These handle making sure that sessions only use the
|
||||||
|
# connections of the appropriate type.
|
||||||
mock_triples += ((cpool.HTTPConnectionPool, '_get_conn',
|
mock_triples += ((cpool.HTTPConnectionPool, '_get_conn',
|
||||||
self._patched_get_conn(cpool.HTTPConnectionPool,
|
self._patched_get_conn(cpool.HTTPConnectionPool,
|
||||||
lambda : cpool.HTTPConnection)),
|
lambda : cpool.HTTPConnection)),
|
||||||
(cpool.HTTPSConnectionPool, '_get_conn',
|
(cpool.HTTPSConnectionPool, '_get_conn',
|
||||||
self._patched_get_conn(cpool.HTTPSConnectionPool,
|
self._patched_get_conn(cpool.HTTPSConnectionPool,
|
||||||
lambda : cpool.HTTPSConnection)))
|
lambda : cpool.HTTPSConnection)),
|
||||||
return self._build_patchers_from_mock_triples(mock_triples)
|
(cpool.HTTPConnectionPool, '_new_conn',
|
||||||
|
self._patched_new_conn(cpool.HTTPConnectionPool,
|
||||||
|
http_connection_remover)),
|
||||||
|
(cpool.HTTPSConnectionPool, '_new_conn',
|
||||||
|
self._patched_new_conn(cpool.HTTPConnectionPool,
|
||||||
|
https_connection_remover)))
|
||||||
|
|
||||||
|
return itertools.chain(self._build_patchers_from_mock_triples(mock_triples),
|
||||||
|
(http_connection_remover, https_connection_remover))
|
||||||
|
|
||||||
def _patched_get_conn(self, connection_pool_class, connection_class_getter):
|
def _patched_get_conn(self, connection_pool_class, connection_class_getter):
|
||||||
get_conn = connection_pool_class._get_conn
|
get_conn = connection_pool_class._get_conn
|
||||||
@@ -153,6 +167,15 @@ class CassettePatcherBuilder(object):
|
|||||||
return connection
|
return connection
|
||||||
return patched_get_conn
|
return patched_get_conn
|
||||||
|
|
||||||
|
def _patched_new_conn(self, connection_pool_class, connection_remover):
|
||||||
|
new_conn = connection_pool_class._new_conn
|
||||||
|
@functools.wraps(new_conn)
|
||||||
|
def patched_new_conn(pool):
|
||||||
|
new_connection = new_conn(pool)
|
||||||
|
connection_remover.add_connection_to_pool_entry(pool, new_connection)
|
||||||
|
return new_connection
|
||||||
|
return patched_new_conn
|
||||||
|
|
||||||
@_build_patchers_from_mock_triples_decorator
|
@_build_patchers_from_mock_triples_decorator
|
||||||
def _urllib3(self):
|
def _urllib3(self):
|
||||||
try:
|
try:
|
||||||
@@ -191,6 +214,36 @@ class CassettePatcherBuilder(object):
|
|||||||
yield cpool, 'CertValidatingHTTPSConnection', VCRCertValidatingHTTPSConnection
|
yield cpool, 'CertValidatingHTTPSConnection', VCRCertValidatingHTTPSConnection
|
||||||
|
|
||||||
|
|
||||||
|
class ConnectionRemover(object):
|
||||||
|
|
||||||
|
def __init__(self, connection_class):
|
||||||
|
self._connection_class = connection_class
|
||||||
|
self._connection_pool_to_connections = {}
|
||||||
|
|
||||||
|
def add_connection_to_pool_entry(self, pool, connection):
|
||||||
|
if isinstance(connection, self._connection_class):
|
||||||
|
self._connection_pool_to_connection.setdefault(pool, set()).add(connection)
|
||||||
|
|
||||||
|
def remove_connection_to_pool_entry(self, pool, connection):
|
||||||
|
if isinstance(connection, self._connection_class):
|
||||||
|
self._connection_pool_to_connection[self._connection_class].remove(connection)
|
||||||
|
|
||||||
|
def __enter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *args):
|
||||||
|
for pool, connections in self._connection_pool_to_connections.items():
|
||||||
|
readd_connections = []
|
||||||
|
while pool.not_empty() and connections:
|
||||||
|
connection = pool.get()
|
||||||
|
if isinstance(connection, self._connection_class):
|
||||||
|
connections.remove(connection)
|
||||||
|
else:
|
||||||
|
readd_connections.append(connection)
|
||||||
|
for connection in readd_connections:
|
||||||
|
self.pool._put_conn(connection)
|
||||||
|
|
||||||
|
|
||||||
def reset_patchers():
|
def reset_patchers():
|
||||||
yield mock.patch.object(httplib, 'HTTPConnection', _HTTPConnection)
|
yield mock.patch.object(httplib, 'HTTPConnection', _HTTPConnection)
|
||||||
yield mock.patch.object(httplib, 'HTTPSConnection', _HTTPSConnection)
|
yield mock.patch.object(httplib, 'HTTPSConnection', _HTTPSConnection)
|
||||||
|
|||||||
Reference in New Issue
Block a user