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