def test_pool_with_framed_protocol_factory(self):
def framed_protocol_factory(trans):
trans = TTransport.TFramedTransport(trans)
return TBinaryProtocol.TBinaryProtocol(trans)
framed_pool = thrift_pool.ThriftConnectionPool(
EXAMPLE_ENDPOINT, protocol_factory=framed_protocol_factory)
trans = thrift_pool._make_transport(EXAMPLE_ENDPOINT)
prot = framed_pool.protocol_factory(trans)
self.assertTrue(isinstance(prot, TBinaryProtocol.TBinaryProtocol))
self.assertTrue(isinstance(prot.trans, TTransport.TFramedTransport))
def test_pool_checkout_exception(self):
pool = thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT)
def mock_get():
raise Exception
pool.pool.get = mock_get
with self.assertRaises(Exception):
with pool.connection() as _:
pass
self.assertEqual(0, pool.checkedout)
def test_pool_checkout_exception(self, mock_queue_module):
# We can't patch gevent queue methods directly as they are implemented in C.
# Instead we patch the whole queue class.
class PatchedLifoQueue(gevent.queue.LifoQueue):
def get(self, *args, **kwargs):
raise Exception
mock_queue_module.LifoQueue = PatchedLifoQueue
pool = thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT)
with self.assertRaises(Exception):
with pool.connection() as _:
pass
self.assertEqual(0, pool.checkedout)
def test_max_connection_attempts_works(self, new_retry_policy):
thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT,
max_connection_attempts=5)
new_retry_policy.assert_called_with(attempts=5)
def test_max_retries_still_works_but_deprecated(self, new_retry_policy):
with pytest.deprecated_call():
thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT, max_retries=5)
new_retry_policy.assert_called_with(attempts=5)
def test_dont_pass_both(self, new_retry_policy):
with self.assertRaises(Exception):
thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT,
max_retries=5,
max_connection_attempts=5)
def test_default_is_3(self, new_retry_policy):
thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT)
new_retry_policy.assert_called_with(attempts=3)
def setUp(self):
self.mock_queue = mock.Mock(spec=queue.Queue)
self.pool = thrift_pool.ThriftConnectionPool(EXAMPLE_ENDPOINT)
self.pool.pool = self.mock_queue
