def test_send_broker_unaware_request(self):
mocked_conns = {
('kafka01', 9092): MagicMock(),
('kafka02', 9092): MagicMock(),
('kafka03', 9092): MagicMock()
}
# inject BrokerConnection side effects
mock_conn(mocked_conns[('kafka01', 9092)], success=False)
mock_conn(mocked_conns[('kafka03', 9092)], success=False)
future = Future()
mocked_conns[('kafka02', 9092)].send.return_value = future
mocked_conns[(
'kafka02',
9092)].recv.side_effect = lambda: future.success('valid response')
def mock_get_conn(host, port, afi, node_id='bootstrap'):
return mocked_conns[(host, port)]
# patch to avoid making requests before we want it
with patch.object(SimpleClient, 'load_metadata_for_topics'):
with patch.object(SimpleClient,
'_get_conn',
side_effect=mock_get_conn):
client = SimpleClient(hosts='kafka01:9092,kafka02:9092')
resp = client._send_broker_unaware_request(
payloads=['fake request'],
encoder_fn=MagicMock(),
decoder_fn=lambda x: x)
self.assertEqual('valid response', resp)
mocked_conns[('kafka02', 9092)].recv.assert_called_once_with()
def test_send_broker_unaware_request_fail(self, load_metadata, conn):
mocked_conns = {
('kafka01', 9092): MagicMock(),
('kafka02', 9092): MagicMock()
}
for val in mocked_conns.values():
mock_conn(val, success=False)
def mock_get_conn(host, port, afi):
return mocked_conns[(host, port)]
conn.side_effect = mock_get_conn
client = SimpleClient(hosts=['kafka01:9092', 'kafka02:9092'])
req = KafkaProtocol.encode_metadata_request()
with self.assertRaises(KafkaUnavailableError):
client._send_broker_unaware_request(payloads=['fake request'],
encoder_fn=MagicMock(return_value='fake encoded message'),
decoder_fn=lambda x: x)
for key, conn in six.iteritems(mocked_conns):
conn.send.assert_called_with('fake encoded message')
