def test_get_leader_for_partitions_reloads_metadata(self, protocol, conn): "Get leader for partitions reload metadata if it is not available" mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [(NO_LEADER, 'topic_no_partitions', [])] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) # topic metadata is loaded but empty self.assertDictEqual({}, client.topics_to_brokers) topics = [(NO_ERROR, 'topic_one_partition', [(NO_ERROR, 0, 0, [0, 1], [0, 1])])] 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 = client._get_leader_for_partition('topic_one_partition', 0) self.assertEqual(brokers[0], leader) self.assertDictEqual( {TopicPartition('topic_one_partition', 0): brokers[0]}, client.topics_to_brokers)
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)
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(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)
def test_get_leader_exceptions_when_noleader(self, protocol, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ (NO_ERROR, 'topic_noleader', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, -1, [], []), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) self.assertDictEqual( { TopicPartition('topic_noleader', 0): None, TopicPartition('topic_noleader', 1): None }, client.topics_to_brokers) # No leader partitions -- raise LeaderNotAvailableError with self.assertRaises(LeaderNotAvailableError): self.assertIsNone( client._get_leader_for_partition('topic_noleader', 0)) with self.assertRaises(LeaderNotAvailableError): self.assertIsNone( client._get_leader_for_partition('topic_noleader', 1)) # Unknown partitions -- raise UnknownTopicOrPartitionError with self.assertRaises(UnknownTopicOrPartitionError): self.assertIsNone( client._get_leader_for_partition('topic_noleader', 2)) topics = [ (NO_ERROR, 'topic_noleader', [(NO_ERROR, 0, 0, [0, 1], [0, 1]), (NO_ERROR, 1, 1, [1, 0], [1, 0])]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) self.assertEqual(brokers[0], client._get_leader_for_partition('topic_noleader', 0)) self.assertEqual(brokers[1], client._get_leader_for_partition('topic_noleader', 1))
async def test_invalid_correlation_id(self): host, port = self.kafka_host, self.kafka_port request = MetadataRequest([]) # setup connection with mocked reader and writer conn = AIOKafkaConnection(host=host, port=port) # setup reader reader = mock.MagicMock() int32 = struct.Struct('>i') resp = MetadataResponse(brokers=[], topics=[]) resp = resp.encode() resp = int32.pack(999) + resp # set invalid correlation id async def first_resp(*args: Any, **kw: Any): return int32.pack(len(resp)) async def second_resp(*args: Any, **kw: Any): return resp reader.readexactly.side_effect = [first_resp(), second_resp()] writer = mock.MagicMock() conn._reader = reader conn._writer = writer # invoke reader task conn._read_task = conn._create_reader_task() with self.assertRaises(CorrelationIdError): await conn.send(request)
def test_invalid_correlation_id(self): host, port = self.kafka_host, self.kafka_port request = MetadataRequest([]) # setup connection with mocked reader and writer conn = AIOKafkaConnection(host=host, port=port, loop=self.loop) # setup reader reader = mock.MagicMock() int32 = struct.Struct('>i') resp = MetadataResponse(brokers=[], topics=[]).encode() resp = int32.pack(999) + resp # set invalid correlation id reader.readexactly.side_effect = [ asyncio.coroutine(lambda *a, **kw: int32.pack(len(resp)))(), asyncio.coroutine(lambda *a, **kw: resp)() ] writer = mock.MagicMock() conn._reader = reader conn._writer = writer # invoke reader task conn._read_task = asyncio. async (conn._read(), loop=self.loop) with self.assertRaises(CorrelationIdError): yield from conn.send(request)
def test_pattern_subscription(coordinator, api_version): coordinator.config['api_version'] = api_version coordinator._subscription.subscribe(pattern='foo') assert coordinator._subscription.subscription == set([]) assert coordinator._subscription_metadata_changed() is False assert coordinator._subscription.needs_partition_assignment is False cluster = coordinator._client.cluster cluster.update_metadata(MetadataResponse( # brokers [(0, 'foo', 12), (1, 'bar', 34)], # topics [(0, 'fizz', []), (0, 'foo1', [(0, 0, 0, [], [])]), (0, 'foo2', [(0, 0, 1, [], [])])])) assert coordinator._subscription.subscription == set(['foo1', 'foo2']) # 0.9 consumers should trigger dynamic partition assignment if api_version >= (0, 9): assert coordinator._subscription.needs_partition_assignment is True assert coordinator._subscription.assignment == {} # earlier consumers get all partitions assigned locally else: assert coordinator._subscription.needs_partition_assignment is False assert set(coordinator._subscription.assignment.keys()) == set([ TopicPartition('foo1', 0), TopicPartition('foo2', 0)])
def test_has_metadata_for_topic(self, protocol, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ (NO_LEADER, 'topic_still_creating', []), (UNKNOWN_TOPIC_OR_PARTITION, 'topic_doesnt_exist', []), (NO_ERROR, 'topic_noleaders', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, -1, [], []), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) # 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_send_produce_request_raises_when_topic_unknown( self, protocol, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ (UNKNOWN_TOPIC_OR_PARTITION, 'topic_doesnt_exist', []), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) requests = [ ProduceRequestPayload( "topic_doesnt_exist", 0, [create_message("a"), create_message("b")]) ] with self.assertRaises(UnknownTopicOrPartitionError): client.send_produce_request(requests)
def test_send_produce_request_raises_when_noleader(self, protocol, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ (NO_ERROR, 'topic_noleader', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, -1, [], []), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) requests = [ ProduceRequestPayload( "topic_noleader", 0, [create_message("a"), create_message("b")]) ] with self.assertRaises(LeaderNotAvailableError): client.send_produce_request(requests)
def test_get_leader_for_unassigned_partitions(self, protocol, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ (NO_LEADER, 'topic_no_partitions', []), (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) self.assertDictEqual({}, client.topics_to_brokers) with self.assertRaises(LeaderNotAvailableError): client._get_leader_for_partition('topic_no_partitions', 0) with self.assertRaises(UnknownTopicOrPartitionError): client._get_leader_for_partition('topic_unknown', 0)
def test_ensure_topic_exists(self, decode_metadata_response, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ (NO_LEADER, 'topic_still_creating', []), (UNKNOWN_TOPIC_OR_PARTITION, 'topic_doesnt_exist', []), (NO_ERROR, 'topic_noleaders', [ (NO_LEADER, 0, -1, [], []), (NO_LEADER, 1, -1, [], []), ]), ] decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = SimpleClient(hosts=['broker_1:4567']) with self.assertRaises(UnknownTopicOrPartitionError): client.ensure_topic_exists('topic_doesnt_exist', timeout=1) with self.assertRaises(KafkaTimeoutError): client.ensure_topic_exists('topic_still_creating', timeout=1) # This should not raise client.ensure_topic_exists('topic_noleaders', timeout=1)
def conn(mocker): conn = mocker.patch('kafka.client_async.BrokerConnection') conn.return_value = conn conn.state = ConnectionStates.CONNECTED conn.send.return_value = Future().success( MetadataResponse( [(0, 'foo', 12), (1, 'bar', 34)], # brokers [])) # topics return conn
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_load_metadata(self, protocol, conn): mock_conn(conn) brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [(NO_ERROR, 'topic_1', [(NO_ERROR, 0, 1, [1, 2], [1, 2])]), (NO_ERROR, 'topic_noleader', [ (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])])] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) # client loads metadata at init client = SimpleClient(hosts=['broker_1:4567']) self.assertDictEqual( { TopicPartition('topic_1', 0): brokers[1], TopicPartition('topic_noleader', 0): None, TopicPartition('topic_noleader', 1): None, TopicPartition('topic_3', 0): brokers[0], TopicPartition('topic_3', 1): brokers[1], TopicPartition('topic_3', 2): brokers[0] }, client.topics_to_brokers) # if we ask for metadata explicitly, it should raise errors with self.assertRaises(LeaderNotAvailableError): client.load_metadata_for_topics('topic_no_partitions') with self.assertRaises(UnknownTopicOrPartitionError): client.load_metadata_for_topics('topic_unknown') # This should not raise client.load_metadata_for_topics('topic_no_leader')
async def send(request_id): return MetadataResponse(brokers, topics)
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'})