mirror of
https://github.com/kevin1024/vcrpy.git
synced 2025-12-08 16:53:23 +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
|
||||
return ()
|
||||
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 = (
|
||||
(cpool, 'VerifiedHTTPSConnection', VCRRequestsHTTPSConnection),
|
||||
(cpool, 'VerifiedHTTPSConnection', VCRRequestsHTTPSConnection),
|
||||
@@ -130,16 +136,24 @@ class CassettePatcherBuilder(object):
|
||||
(cpool, 'HTTPSConnection', VCRRequestsHTTPSConnection),
|
||||
(cpool.HTTPConnectionPool, 'ConnectionCls', VCRRequestsHTTPConnection),
|
||||
(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',
|
||||
self._patched_get_conn(cpool.HTTPConnectionPool,
|
||||
lambda : cpool.HTTPConnection)),
|
||||
(cpool.HTTPSConnectionPool, '_get_conn',
|
||||
self._patched_get_conn(cpool.HTTPSConnectionPool,
|
||||
lambda : cpool.HTTPSConnection)))
|
||||
return self._build_patchers_from_mock_triples(mock_triples)
|
||||
lambda : cpool.HTTPSConnection)),
|
||||
(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):
|
||||
get_conn = connection_pool_class._get_conn
|
||||
@@ -153,6 +167,15 @@ class CassettePatcherBuilder(object):
|
||||
return connection
|
||||
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
|
||||
def _urllib3(self):
|
||||
try:
|
||||
@@ -191,6 +214,36 @@ class CassettePatcherBuilder(object):
|
||||
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():
|
||||
yield mock.patch.object(httplib, 'HTTPConnection', _HTTPConnection)
|
||||
yield mock.patch.object(httplib, 'HTTPSConnection', _HTTPSConnection)
|
||||
|
||||
Reference in New Issue
Block a user