def test_has_metadata_for_topic(self, protocol): @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [ TopicMetadata("topic_still_creating", NO_LEADER, []), TopicMetadata("topic_doesnt_exist", UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata( "topic_noleaders", NO_ERROR, [ PartitionMetadata("topic_noleaders", 0, -1, [], [], NO_LEADER), PartitionMetadata("topic_noleaders", 1, -1, [], [], NO_LEADER), ], ), ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) # Topics with no partitions return False self.assertFalse(client.has_metadata_for_topic("topic_still_creating")) self.assertFalse(client.has_metadata_for_topic("topic_doesnt_exist")) # Topic with partition metadata, but no leaders return True self.assertTrue(client.has_metadata_for_topic("topic_noleaders"))
def test_ensure_topic_exists(self, protocol): mocked_conns = {("broker_1", 4567): mock.MagicMock()} fut = asyncio.Future(loop=self.loop) fut.set_result(b"response") mocked_conns[("broker_1", 4567)].send.return_value = fut client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [ TopicMetadata("topic_still_creating", NO_LEADER, []), TopicMetadata("topic_doesnt_exist", UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata( "topic_noleaders", NO_ERROR, [ PartitionMetadata("topic_noleaders", 0, -1, [], [], NO_LEADER), PartitionMetadata("topic_noleaders", 1, -1, [], [], NO_LEADER), ], ), ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete(client.ensure_topic_exists("topic_doesnt_exist", timeout=1)) with self.assertRaises(KafkaTimeoutError): self.loop.run_until_complete(client.ensure_topic_exists("topic_still_creating", timeout=1)) # This should not raise self.loop.run_until_complete(client.ensure_topic_exists("topic_noleaders", timeout=1))
async def test_client_receive_zero_brokers_timeout_on_send(self): brokers = [(0, 'broker_1', 4567), (1, 'broker_2', 5678)] correct_meta = MetadataResponse(brokers, []) async def send(*args, **kwargs): raise asyncio.TimeoutError() client = AIOKafkaClient(bootstrap_servers=['broker_1:4567'], api_version="0.10") conn = mock.Mock() client._conns = [mock.Mock()] async def _get_conn(*args: Any, **kwargs: Any): return conn client._get_conn = mock.Mock() client._get_conn.side_effect = _get_conn conn.send = mock.Mock() conn.send.side_effect = send client.cluster.update_metadata(correct_meta) brokers_before = client.cluster.brokers() await client._metadata_update(client.cluster, []) # There broker list should not be purged self.assertNotEqual(client.cluster.brokers(), set([])) self.assertEqual(client.cluster.brokers(), brokers_before)
async def test_client_receive_zero_brokers(self): brokers = [(0, 'broker_1', 4567), (1, 'broker_2', 5678)] correct_meta = MetadataResponse(brokers, []) bad_response = MetadataResponse([], []) async def send(*args, **kwargs): return bad_response client = AIOKafkaClient(loop=self.loop, bootstrap_servers=['broker_1:4567'], api_version="0.10") conn = mock.Mock() client._conns = [mock.Mock()] client._get_conn = mock.Mock() client._get_conn.side_effect = asyncio.coroutine(lambda x: conn) conn.send = mock.Mock() conn.send.side_effect = send client.cluster.update_metadata(correct_meta) brokers_before = client.cluster.brokers() await client._metadata_update(client.cluster, []) # There broker list should not be purged self.assertNotEqual(client.cluster.brokers(), set([])) self.assertEqual(client.cluster.brokers(), brokers_before)
def test_send_produce_request_raises_when_topic_unknown(self, protocol): @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_doesnt_exist', UNKNOWN_TOPIC_OR_PARTITION, []), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) requests = [ ProduceRequest( "topic_doesnt_exist", 0, [create_message("a"), create_message("b")]) ] with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete(client.send_produce_request(requests))
def test_send_broker_unaware_request_fail(self): 'Tests that call fails when all hosts are unavailable' mocked_conns = { ('kafka01', 9092): mock.MagicMock(), ('kafka02', 9092): mock.MagicMock() } # inject KafkaConnection side effects fut1 = asyncio.Future(loop=self.loop) fut1.set_exception(RuntimeError("kafka01 went away (unittest)")) mocked_conns[('kafka01', 9092)].send.return_value = fut1 fut2 = asyncio.Future(loop=self.loop) fut2.set_exception(RuntimeError("kafka02 went away (unittest)")) mocked_conns[('kafka02', 9092)].send.return_value = fut2 client = AIOKafkaClient(['kafka01:9092', 'kafka02:9092'], loop=self.loop) client._conns = mocked_conns @asyncio.coroutine def go(): with self.assertRaises(KafkaUnavailableError): yield from client._send_broker_unaware_request( payloads=['fake request'], encoder_fn=mock.MagicMock( return_value=b'fake encoded message'), decoder_fn=lambda x: x) for key, conn in mocked_conns.items(): conn.send.assert_called_with(b'fake encoded message') self.loop.run_until_complete(go())
def test_has_metadata_for_topic(self, protocol): @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_still_creating', NO_LEADER, []), TopicMetadata('topic_doesnt_exist', UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata('topic_noleaders', NO_ERROR, [ PartitionMetadata('topic_noleaders', 0, -1, [], [], NO_LEADER), PartitionMetadata('topic_noleaders', 1, -1, [], [], NO_LEADER), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) # Topics with no partitions return False self.assertFalse(client.has_metadata_for_topic('topic_still_creating')) self.assertFalse(client.has_metadata_for_topic('topic_doesnt_exist')) # Topic with partition metadata, but no leaders return True self.assertTrue(client.has_metadata_for_topic('topic_noleaders'))
def test_client_receive_zero_brokers(self): brokers = [ (0, 'broker_1', 4567), (1, 'broker_2', 5678) ] correct_meta = MetadataResponse(brokers, []) bad_response = MetadataResponse([], []) @asyncio.coroutine def send(*args, **kwargs): return bad_response client = AIOKafkaClient(loop=self.loop, bootstrap_servers=['broker_1:4567'], api_version="0.10") conn = mock.Mock() client._conns = [mock.Mock()] client._get_conn = mock.Mock() client._get_conn.side_effect = asyncio.coroutine(lambda x: conn) conn.send = mock.Mock() conn.send.side_effect = send client.cluster.update_metadata(correct_meta) brokers_before = client.cluster.brokers() yield from client._metadata_update(client.cluster, []) # There broker list should not be purged self.assertNotEqual(client.cluster.brokers(), set([])) self.assertEqual(client.cluster.brokers(), brokers_before)
def test_send_broker_unaware_request(self): "Tests that call works when at least one of the host is available" mocked_conns = { ("kafka01", 9092): mock.MagicMock(), ("kafka02", 9092): mock.MagicMock(), ("kafka03", 9092): mock.MagicMock(), } # inject KafkaConnection side effects fut = asyncio.Future(loop=self.loop) fut.set_exception(RuntimeError("kafka01 went away (unittest)")) mocked_conns[("kafka01", 9092)].send.return_value = fut fut2 = asyncio.Future(loop=self.loop) fut2.set_result(b"valid response") mocked_conns[("kafka02", 9092)].send.return_value = fut2 fut3 = asyncio.Future(loop=self.loop) fut3.set_exception(RuntimeError("kafka03 went away (unittest)")) mocked_conns[("kafka03", 9092)].send.return_value = fut3 client = AIOKafkaClient("kafka01:9092,kafka02:9092", loop=self.loop) client._conns = mocked_conns resp = self.loop.run_until_complete( client._send_broker_unaware_request( payloads=[b"fake request"], encoder_fn=mock.MagicMock(), decoder_fn=lambda x: x ) ) self.assertEqual(b"valid response", resp)
def test_send_broker_unaware_request(self): 'Tests that call works when at least one of the host is available' mocked_conns = { ('kafka01', 9092): mock.MagicMock(), ('kafka02', 9092): mock.MagicMock(), ('kafka03', 9092): mock.MagicMock() } # inject KafkaConnection side effects fut = asyncio.Future(loop=self.loop) fut.set_exception(RuntimeError("kafka01 went away (unittest)")) mocked_conns[('kafka01', 9092)].send.return_value = fut fut2 = asyncio.Future(loop=self.loop) fut2.set_result(b'valid response') mocked_conns[('kafka02', 9092)].send.return_value = fut2 fut3 = asyncio.Future(loop=self.loop) fut3.set_exception(RuntimeError("kafka03 went away (unittest)")) mocked_conns[('kafka03', 9092)].send.return_value = fut3 client = AIOKafkaClient('kafka01:9092,kafka02:9092', loop=self.loop) client._conns = mocked_conns resp = self.loop.run_until_complete( client._send_broker_unaware_request(payloads=[b'fake request'], encoder_fn=mock.MagicMock(), decoder_fn=lambda x: x)) self.assertEqual(b'valid response', resp)
def test_get_leader_for_unassigned_partitions(self, protocol): @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_no_partitions', NO_LEADER, []), TopicMetadata('topic_unknown', UNKNOWN_TOPIC_OR_PARTITION, []), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) self.assertDictEqual({}, client._topics_to_brokers) with self.assertRaises(LeaderNotAvailableError): self.loop.run_until_complete( client._get_leader_for_partition('topic_no_partitions', 0)) with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete( client._get_leader_for_partition('topic_unknown', 0))
def test_send_broker_unaware_request_fail(self): "Tests that call fails when all hosts are unavailable" mocked_conns = {("kafka01", 9092): mock.MagicMock(), ("kafka02", 9092): mock.MagicMock()} # inject KafkaConnection side effects fut1 = asyncio.Future(loop=self.loop) fut1.set_exception(RuntimeError("kafka01 went away (unittest)")) mocked_conns[("kafka01", 9092)].send.return_value = fut1 fut2 = asyncio.Future(loop=self.loop) fut2.set_exception(RuntimeError("kafka02 went away (unittest)")) mocked_conns[("kafka02", 9092)].send.return_value = fut2 client = AIOKafkaClient(["kafka01:9092", "kafka02:9092"], loop=self.loop) client._conns = mocked_conns @asyncio.coroutine def go(): with self.assertRaises(KafkaUnavailableError): yield from client._send_broker_unaware_request( payloads=["fake request"], encoder_fn=mock.MagicMock(return_value=b"fake encoded message"), decoder_fn=lambda x: x, ) for key, conn in mocked_conns.items(): conn.send.assert_called_with(b"fake encoded message") self.loop.run_until_complete(go())
def test_send_produce_request_raises_when_noleader(self, protocol): """Send producer request raises LeaderNotAvailableError if leader is not available""" @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [ TopicMetadata( "topic_noleader", NO_ERROR, [ PartitionMetadata("topic_noleader", 0, -1, [], [], NO_LEADER), PartitionMetadata("topic_noleader", 1, -1, [], [], NO_LEADER), ], ) ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) requests = [ProduceRequest("topic_noleader", 0, [create_message("a"), create_message("b")])] with self.assertRaises(LeaderNotAvailableError): self.loop.run_until_complete(client.send_produce_request(requests))
def test_get_leader_for_unassigned_partitions(self, protocol): @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [ TopicMetadata("topic_no_partitions", NO_LEADER, []), TopicMetadata("topic_unknown", UNKNOWN_TOPIC_OR_PARTITION, []), ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) self.assertDictEqual({}, client._topics_to_brokers) with self.assertRaises(LeaderNotAvailableError): self.loop.run_until_complete(client._get_leader_for_partition("topic_no_partitions", 0)) with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete(client._get_leader_for_partition("topic_unknown", 0))
def test_close(self): client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) m1 = mock.Mock() m2 = mock.Mock() client._conns = {('host1', 4567): m1, ('host2', 5678): m2} client.close() self.assertEqual({}, client._conns) m1.close.assert_raises_with() m2.close.assert_raises_with()
def test_close(self): client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) m1 = mock.Mock() m2 = mock.Mock() client._conns = {("host1", 4567): m1, ("host2", 5678): m2} client.close() self.assertEqual({}, client._conns) m1.close.assert_raises_with() m2.close.assert_raises_with()
def test_get_leader_exceptions_when_noleader(self, protocol): @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [ TopicMetadata( "topic_noleader", NO_ERROR, [ PartitionMetadata("topic_noleader", 0, -1, [], [], NO_LEADER), PartitionMetadata("topic_noleader", 1, -1, [], [], NO_LEADER), ], ) ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) self.assertDictEqual( {TopicAndPartition("topic_noleader", 0): None, TopicAndPartition("topic_noleader", 1): None}, client._topics_to_brokers, ) # No leader partitions -- raise LeaderNotAvailableError with self.assertRaises(LeaderNotAvailableError): self.assertIsNone(self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 0))) with self.assertRaises(LeaderNotAvailableError): self.assertIsNone(self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 1))) # Unknown partitions -- raise UnknownTopicOrPartitionError with self.assertRaises(UnknownTopicOrPartitionError): self.assertIsNone(self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 2))) topics = [ TopicMetadata( "topic_noleader", NO_ERROR, [ PartitionMetadata("topic_noleader", 0, 0, [0, 1], [0, 1], NO_ERROR), PartitionMetadata("topic_noleader", 1, 1, [1, 0], [1, 0], NO_ERROR), ], ) ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.assertEqual( brokers[0], self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 0)) ) self.assertEqual( brokers[1], self.loop.run_until_complete(client._get_leader_for_partition("topic_noleader", 1)) )
def test_load_metadata(self): brokers = [(0, 'broker_1', 4567), (1, 'broker_2', 5678)] topics = [(NO_ERROR, 'topic_1', [(NO_ERROR, 0, 1, [1, 2], [1, 2])]), (NO_ERROR, 'topic_2', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, 1, [], []), ]), (NO_LEADER, 'topic_no_partitions', []), (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []), (NO_ERROR, 'topic_3', [(NO_ERROR, 0, 0, [0, 1], [0, 1]), (NO_ERROR, 1, 1, [1, 0], [1, 0]), (NO_ERROR, 2, 0, [0, 1], [0, 1])]), (NO_ERROR, 'topic_4', [ (NO_ERROR, 0, 0, [0, 1], [0, 1]), (REPLICA_NOT_AVAILABLE, 1, 1, [1, 0], [1, 0]), ])] @asyncio.coroutine def send(request_id): return MetadataResponse(brokers, topics) mocked_conns = {0: mock.MagicMock()} mocked_conns[0].send.side_effect = send client = AIOKafkaClient(loop=self.loop, bootstrap_servers=['broker_1:4567']) task = asyncio. async (client._md_synchronizer(), loop=self.loop) client._conns = mocked_conns client.cluster.update_metadata(MetadataResponse(brokers[:1], [])) self.loop.run_until_complete(client.force_metadata_update()) task.cancel() md = client.cluster c_brokers = md.brokers() self.assertEqual(len(c_brokers), 2) expected_brokers = [(0, 'broker_1', 4567, None), (1, 'broker_2', 5678, None)] self.assertEqual(sorted(expected_brokers), sorted(list(c_brokers))) c_topics = md.topics() self.assertEqual(len(c_topics), 4) self.assertEqual(md.partitions_for_topic('topic_1'), set([0])) self.assertEqual(md.partitions_for_topic('topic_2'), set([0, 1])) self.assertEqual(md.partitions_for_topic('topic_3'), set([0, 1, 2])) self.assertEqual(md.partitions_for_topic('topic_4'), set([0, 1])) self.assertEqual(md.available_partitions_for_topic('topic_2'), set([1])) mocked_conns[0].connected.return_value = False is_ready = self.loop.run_until_complete(client.ready(0)) self.assertEqual(is_ready, False) is_ready = self.loop.run_until_complete(client.ready(1)) self.assertEqual(is_ready, False) self.assertEqual(mocked_conns, {}) with self.assertRaises(NodeNotReadyError): self.loop.run_until_complete(client.send(0, None))
def test_send_timeout_deletes_connection(self): correct_response = MetadataResponse([], []) @asyncio.coroutine def send_exception(*args, **kwargs): raise asyncio.TimeoutError() @asyncio.coroutine def send(*args, **kwargs): return correct_response @asyncio.coroutine def get_conn(self, node_id, *, group=0): conn_id = (node_id, group) if conn_id in self._conns: conn = self._conns[conn_id] if not conn.connected(): del self._conns[conn_id] else: return conn conn = mock.MagicMock() conn.send.side_effect = send self._conns[conn_id] = conn return conn node_id = 0 conn = mock.MagicMock() conn.send.side_effect = send_exception conn.connected.return_value = True mocked_conns = {(node_id, 0): conn} client = AIOKafkaClient(loop=self.loop, bootstrap_servers=['broker_1:4567']) client._conns = mocked_conns client._get_conn = types.MethodType(get_conn, client) # first send timeouts with self.assertRaises(RequestTimedOutError): yield from client.send(0, MetadataRequest([])) conn.close.assert_called_once_with( reason=CloseReason.CONNECTION_TIMEOUT) # this happens because conn was closed conn.connected.return_value = False # second send gets new connection and obtains result response = yield from client.send(0, MetadataRequest([])) self.assertEqual(response, correct_response) self.assertNotEqual(conn, client._conns[(node_id, 0)])
def test_get_leader_for_partitions_reloads_metadata(self, protocol): "Get leader for partitions reload metadata if it is not available" @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [TopicMetadata('topic_no_partitions', NO_LEADER, [])] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) # topic metadata is loaded but empty self.assertDictEqual({}, client._topics_to_brokers) topics = [ TopicMetadata('topic_one_partition', NO_ERROR, [ PartitionMetadata('topic_no_partition', 0, 0, [0, 1], [0, 1], NO_ERROR) ]) ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) # calling _get_leader_for_partition (from any broker aware request) # will try loading metadata again for the same topic leader = self.loop.run_until_complete( client._get_leader_for_partition('topic_one_partition', 0)) self.assertEqual(brokers[0], leader) self.assertDictEqual( {TopicAndPartition('topic_one_partition', 0): brokers[0]}, client._topics_to_brokers)
def test_send_produce_request_raises_when_topic_unknown(self, protocol): @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [TopicMetadata("topic_doesnt_exist", UNKNOWN_TOPIC_OR_PARTITION, [])] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) requests = [ProduceRequest("topic_doesnt_exist", 0, [create_message("a"), create_message("b")])] with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete(client.send_produce_request(requests))
def test_ensure_topic_exists(self, protocol): mocked_conns = {('broker_1', 4567): mock.MagicMock()} fut = asyncio.Future(loop=self.loop) fut.set_result(b'response') mocked_conns[('broker_1', 4567)].send.return_value = fut client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_still_creating', NO_LEADER, []), TopicMetadata('topic_doesnt_exist', UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata('topic_noleaders', NO_ERROR, [ PartitionMetadata('topic_noleaders', 0, -1, [], [], NO_LEADER), PartitionMetadata('topic_noleaders', 1, -1, [], [], NO_LEADER), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete( client.ensure_topic_exists('topic_doesnt_exist', timeout=1)) with self.assertRaises(KafkaTimeoutError): self.loop.run_until_complete( client.ensure_topic_exists('topic_still_creating', timeout=1)) # This should not raise self.loop.run_until_complete( client.ensure_topic_exists('topic_noleaders', timeout=1))
def test_get_leader_for_partitions_reloads_metadata(self, protocol): "Get leader for partitions reload metadata if it is not available" @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [TopicMetadata("topic_no_partitions", NO_LEADER, [])] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) # topic metadata is loaded but empty self.assertDictEqual({}, client._topics_to_brokers) topics = [ TopicMetadata( "topic_one_partition", NO_ERROR, [PartitionMetadata("topic_no_partition", 0, 0, [0, 1], [0, 1], NO_ERROR)], ) ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) # calling _get_leader_for_partition (from any broker aware request) # will try loading metadata again for the same topic leader = self.loop.run_until_complete(client._get_leader_for_partition("topic_one_partition", 0)) self.assertEqual(brokers[0], leader) self.assertDictEqual({TopicAndPartition("topic_one_partition", 0): brokers[0]}, client._topics_to_brokers)
def test_send_produce_request_raises_when_noleader(self, protocol): """Send producer request raises LeaderNotAvailableError if leader is not available""" @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_noleader', NO_ERROR, [ PartitionMetadata('topic_noleader', 0, -1, [], [], NO_LEADER), PartitionMetadata('topic_noleader', 1, -1, [], [], NO_LEADER), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) requests = [ ProduceRequest( "topic_noleader", 0, [create_message("a"), create_message("b")]) ] with self.assertRaises(LeaderNotAvailableError): self.loop.run_until_complete(client.send_produce_request(requests))
async def test_load_metadata(self): brokers = [(0, 'broker_1', 4567), (1, 'broker_2', 5678)] topics = [ (NO_ERROR, 'topic_1', [(NO_ERROR, 0, 1, [1, 2], [1, 2])]), (NO_ERROR, 'topic_2', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, 1, [], []), ]), (NO_LEADER, 'topic_no_partitions', []), (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []), (NO_ERROR, 'topic_3', [(NO_ERROR, 0, 0, [0, 1], [0, 1]), (NO_ERROR, 1, 1, [1, 0], [1, 0]), (NO_ERROR, 2, 0, [0, 1], [0, 1])]), (NO_ERROR, 'topic_4', [ (NO_ERROR, 0, 0, [0, 1], [0, 1]), (REPLICA_NOT_AVAILABLE, 1, 1, [1, 0], [1, 0]), ]), (INVALID_TOPIC, 'topic_5', []), # Just ignored (UNKNOWN_ERROR, 'topic_6', []), # Just ignored (TOPIC_AUTHORIZATION_FAILED, 'topic_auth_error', []), ] async def send(request_id): return MetadataResponse(brokers, topics) mocked_conns = {(0, 0): mock.MagicMock()} mocked_conns[(0, 0)].send.side_effect = send client = AIOKafkaClient(bootstrap_servers=['broker_1:4567']) task = create_task(client._md_synchronizer()) client._conns = mocked_conns client.cluster.update_metadata(MetadataResponse(brokers[:1], [])) await client.force_metadata_update() task.cancel() md = client.cluster c_brokers = md.brokers() self.assertEqual(len(c_brokers), 2) expected_brokers = [(0, 'broker_1', 4567, None), (1, 'broker_2', 5678, None)] self.assertEqual(sorted(expected_brokers), sorted(list(c_brokers))) c_topics = md.topics() self.assertEqual(len(c_topics), 4) self.assertEqual(md.partitions_for_topic('topic_1'), set([0])) self.assertEqual(md.partitions_for_topic('topic_2'), set([0, 1])) self.assertEqual(md.partitions_for_topic('topic_3'), set([0, 1, 2])) self.assertEqual(md.partitions_for_topic('topic_4'), set([0, 1])) self.assertEqual(md.available_partitions_for_topic('topic_2'), set([1])) mocked_conns[(0, 0)].connected.return_value = False is_ready = await client.ready(0) self.assertEqual(is_ready, False) is_ready = await client.ready(1) self.assertEqual(is_ready, False) self.assertEqual(mocked_conns, {}) with self.assertRaises(NodeNotReadyError): await client.send(0, None) self.assertEqual(md.unauthorized_topics, {'topic_auth_error'})
def test_load_metadata(self, protocol): @asyncio.coroutine def recv(request_id): return b"response" mocked_conns = {("broker_1", 4567): mock.MagicMock()} mocked_conns[("broker_1", 4567)].recv.side_effect = recv client = AIOKafkaClient(["broker_1:4567"], loop=self.loop) client._conns = mocked_conns brokers = [BrokerMetadata(0, "broker_1", 4567), BrokerMetadata(1, "broker_2", 5678)] topics = [ TopicMetadata("topic_1", NO_ERROR, [PartitionMetadata("topic_1", 0, 1, [1, 2], [1, 2], NO_ERROR)]), TopicMetadata( "topic_noleader", NO_ERROR, [ PartitionMetadata("topic_noleader", 0, -1, [], [], NO_LEADER), PartitionMetadata("topic_noleader", 1, -1, [], [], NO_LEADER), ], ), TopicMetadata("topic_no_partitions", NO_LEADER, []), TopicMetadata("topic_unknown", UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata( "topic_3", NO_ERROR, [ PartitionMetadata("topic_3", 0, 0, [0, 1], [0, 1], NO_ERROR), PartitionMetadata("topic_3", 1, 1, [1, 0], [1, 0], NO_ERROR), PartitionMetadata("topic_3", 2, 0, [0, 1], [0, 1], NO_ERROR), ], ), TopicMetadata( "topic_4", NO_ERROR, [ PartitionMetadata("topic_4", 0, 0, [0, 1], [0, 1], NO_ERROR), PartitionMetadata("topic_4", 1, 1, [1, 0], [1, 0], REPLICA_NOT_AVAILABLE), ], ), ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) self.assertDictEqual( { TopicAndPartition("topic_1", 0): brokers[1], TopicAndPartition("topic_noleader", 0): None, TopicAndPartition("topic_noleader", 1): None, TopicAndPartition("topic_3", 0): brokers[0], TopicAndPartition("topic_3", 1): brokers[1], TopicAndPartition("topic_3", 2): brokers[0], TopicAndPartition("topic_4", 0): brokers[0], TopicAndPartition("topic_4", 1): brokers[1], }, client._topics_to_brokers, ) # if we ask for metadata explicitly, it should raise errors with self.assertRaises(LeaderNotAvailableError): self.loop.run_until_complete(client.load_metadata_for_topics("topic_no_partitions")) with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete(client.load_metadata_for_topics("topic_unknown")) # This should not raise self.loop.run_until_complete(client.load_metadata_for_topics("topic_no_leader")) # This should not raise ReplicaNotAvailableError self.loop.run_until_complete(client.load_metadata_for_topics("topic_4"))
def test_load_metadata(self): brokers = [ (0, 'broker_1', 4567), (1, 'broker_2', 5678) ] topics = [ (NO_ERROR, 'topic_1', [ (NO_ERROR, 0, 1, [1, 2], [1, 2]) ]), (NO_ERROR, 'topic_2', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, 1, [], []), ]), (NO_LEADER, 'topic_no_partitions', []), (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []), (NO_ERROR, 'topic_3', [ (NO_ERROR, 0, 0, [0, 1], [0, 1]), (NO_ERROR, 1, 1, [1, 0], [1, 0]), (NO_ERROR, 2, 0, [0, 1], [0, 1]) ]), (NO_ERROR, 'topic_4', [ (NO_ERROR, 0, 0, [0, 1], [0, 1]), (REPLICA_NOT_AVAILABLE, 1, 1, [1, 0], [1, 0]), ]) ] @asyncio.coroutine def send(request_id): return MetadataResponse(brokers, topics) mocked_conns = {0: mock.MagicMock()} mocked_conns[0].send.side_effect = send client = AIOKafkaClient(loop=self.loop, bootstrap_servers=['broker_1:4567']) task = asyncio.async(client._md_synchronizer(), loop=self.loop) client._conns = mocked_conns client.cluster.update_metadata(MetadataResponse(brokers[:1], [])) self.loop.run_until_complete(client.force_metadata_update()) task.cancel() md = client.cluster c_brokers = md.brokers() self.assertEqual(len(c_brokers), 2) self.assertEqual(sorted(brokers), sorted(list(c_brokers))) c_topics = md.topics() self.assertEqual(len(c_topics), 4) self.assertEqual(md.partitions_for_topic('topic_1'), set([0])) self.assertEqual(md.partitions_for_topic('topic_2'), set([0, 1])) self.assertEqual(md.partitions_for_topic('topic_3'), set([0, 1, 2])) self.assertEqual(md.partitions_for_topic('topic_4'), set([0, 1])) self.assertEqual( md.available_partitions_for_topic('topic_2'), set([1])) mocked_conns[0].connected.return_value = False is_ready = self.loop.run_until_complete(client.ready(0)) self.assertEqual(is_ready, False) is_ready = self.loop.run_until_complete(client.ready(1)) self.assertEqual(is_ready, False) self.assertEqual(mocked_conns, {}) with self.assertRaises(NodeNotReadyError): self.loop.run_until_complete(client.send(0, None))
def test_load_metadata(self, protocol): @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata( 'topic_1', NO_ERROR, [PartitionMetadata('topic_1', 0, 1, [1, 2], [1, 2], NO_ERROR) ]), TopicMetadata('topic_noleader', NO_ERROR, [ PartitionMetadata('topic_noleader', 0, -1, [], [], NO_LEADER), PartitionMetadata('topic_noleader', 1, -1, [], [], NO_LEADER), ]), TopicMetadata('topic_no_partitions', NO_LEADER, []), TopicMetadata('topic_unknown', UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata('topic_3', NO_ERROR, [ PartitionMetadata('topic_3', 0, 0, [0, 1], [0, 1], NO_ERROR), PartitionMetadata('topic_3', 1, 1, [1, 0], [1, 0], NO_ERROR), PartitionMetadata('topic_3', 2, 0, [0, 1], [0, 1], NO_ERROR) ]), TopicMetadata('topic_4', NO_ERROR, [ PartitionMetadata('topic_4', 0, 0, [0, 1], [0, 1], NO_ERROR), PartitionMetadata('topic_4', 1, 1, [1, 0], [1, 0], REPLICA_NOT_AVAILABLE), ]) ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) self.assertDictEqual( { TopicAndPartition('topic_1', 0): brokers[1], TopicAndPartition('topic_noleader', 0): None, TopicAndPartition('topic_noleader', 1): None, TopicAndPartition('topic_3', 0): brokers[0], TopicAndPartition('topic_3', 1): brokers[1], TopicAndPartition('topic_3', 2): brokers[0], TopicAndPartition('topic_4', 0): brokers[0], TopicAndPartition('topic_4', 1): brokers[1] }, client._topics_to_brokers) # if we ask for metadata explicitly, it should raise errors with self.assertRaises(LeaderNotAvailableError): self.loop.run_until_complete( client.load_metadata_for_topics('topic_no_partitions')) with self.assertRaises(UnknownTopicOrPartitionError): self.loop.run_until_complete( client.load_metadata_for_topics('topic_unknown')) # This should not raise self.loop.run_until_complete( client.load_metadata_for_topics('topic_no_leader')) # This should not raise ReplicaNotAvailableError self.loop.run_until_complete( client.load_metadata_for_topics('topic_4'))
def test_get_leader_exceptions_when_noleader(self, protocol): @asyncio.coroutine def recv(request_id): return b'response' mocked_conns = {('broker_1', 4567): mock.MagicMock()} mocked_conns[('broker_1', 4567)].recv.side_effect = recv client = AIOKafkaClient(['broker_1:4567'], loop=self.loop) client._conns = mocked_conns brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_noleader', NO_ERROR, [ PartitionMetadata('topic_noleader', 0, -1, [], [], NO_LEADER), PartitionMetadata('topic_noleader', 1, -1, [], [], NO_LEADER), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.loop.run_until_complete(client.load_metadata_for_topics()) self.assertDictEqual( { TopicAndPartition('topic_noleader', 0): None, TopicAndPartition('topic_noleader', 1): None }, client._topics_to_brokers) # No leader partitions -- raise LeaderNotAvailableError with self.assertRaises(LeaderNotAvailableError): self.assertIsNone( self.loop.run_until_complete( client._get_leader_for_partition('topic_noleader', 0))) with self.assertRaises(LeaderNotAvailableError): self.assertIsNone( self.loop.run_until_complete( client._get_leader_for_partition('topic_noleader', 1))) # Unknown partitions -- raise UnknownTopicOrPartitionError with self.assertRaises(UnknownTopicOrPartitionError): self.assertIsNone( self.loop.run_until_complete( client._get_leader_for_partition('topic_noleader', 2))) topics = [ TopicMetadata('topic_noleader', NO_ERROR, [ PartitionMetadata('topic_noleader', 0, 0, [0, 1], [0, 1], NO_ERROR), PartitionMetadata('topic_noleader', 1, 1, [1, 0], [1, 0], NO_ERROR) ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.assertEqual( brokers[0], self.loop.run_until_complete( client._get_leader_for_partition('topic_noleader', 0))) self.assertEqual( brokers[1], self.loop.run_until_complete( client._get_leader_for_partition('topic_noleader', 1)))