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