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')