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))
async def test_compacted_topic_consumption(self):
# Compacted topics can have offsets skipped
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
subscriptions = SubscriptionState(loop=self.loop)
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp = TopicPartition('test', 0)
req = FetchRequest(
-1, # replica_id
100, 100, [(tp.topic, [(tp.partition, 155, 100000)])])
builder = LegacyRecordBatchBuilder(
magic=1, compression_type=0, batch_size=99999999)
builder.append(160, value=b"12345", key=b"1", timestamp=None)
builder.append(162, value=b"23456", key=b"2", timestamp=None)
builder.append(167, value=b"34567", key=b"3", timestamp=None)
batch = bytes(builder.build())
resp = FetchResponse(
[('test', [(
0, 0, 3000, # partition, error_code, highwater_offset
batch # Batch raw bytes
)])])
subscriptions.assign_from_user({tp})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(tp)
client.send.side_effect = asyncio.coroutine(lambda n, r: resp)
tp_state.seek(155)
fetcher._in_flight.add(0)
needs_wake_up = await fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
# Test successful getone, the closest in batch offset=160
first = buf.getone()
self.assertEqual(tp_state.position, 161)
self.assertEqual(
(first.value, first.key, first.offset),
(b"12345", b"1", 160))
# Test successful getmany
second, third = buf.getall()
self.assertEqual(tp_state.position, 168)
self.assertEqual(
(second.value, second.key, second.offset),
(b"23456", b"2", 162))
self.assertEqual(
(third.value, third.key, third.offset),
(b"34567", b"3", 167))
def test_compacted_topic_consumption(self):
# Compacted topics can have offsets skipped
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
subscriptions = SubscriptionState(loop=self.loop)
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp = TopicPartition('test', 0)
req = FetchRequest(
-1, # replica_id
100, 100, [(tp.topic, [(tp.partition, 155, 100000)])])
builder = LegacyRecordBatchBuilder(
magic=1, compression_type=0, batch_size=99999999)
builder.append(160, value=b"12345", key=b"1", timestamp=None)
builder.append(162, value=b"23456", key=b"2", timestamp=None)
builder.append(167, value=b"34567", key=b"3", timestamp=None)
batch = bytes(builder.build())
resp = FetchResponse(
[('test', [(
0, 0, 3000, # partition, error_code, highwater_offset
batch # Batch raw bytes
)])])
subscriptions.assign_from_user({tp})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(tp)
client.send.side_effect = asyncio.coroutine(lambda n, r: resp)
tp_state.seek(155)
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
# Test successful getone, the closest in batch offset=160
first = buf.getone()
self.assertEqual(tp_state.position, 161)
self.assertEqual(
(first.value, first.key, first.offset),
(b"12345", b"1", 160))
# Test successful getmany
second, third = buf.getall()
self.assertEqual(tp_state.position, 168)
self.assertEqual(
(second.value, second.key, second.offset),
(b"23456", b"2", 162))
self.assertEqual(
(third.value, third.key, third.offset),
(b"34567", b"3", 167))
def test_update_fetch_positions(self):
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
subscriptions = SubscriptionState('latest')
fetcher = Fetcher(client, subscriptions, loop=self.loop)
partition = TopicPartition('test', 0)
# partition is not assigned, should be ignored
yield from fetcher.update_fetch_positions([partition])
state = TopicPartitionState()
state.seek(0)
subscriptions.assignment[partition] = state
# partition is fetchable, no need to update position
yield from fetcher.update_fetch_positions([partition])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
client.send.side_effect = asyncio.coroutine(
lambda n, r: OffsetResponse([('test', [(0, 0, [4])])]))
state.await_reset(OffsetResetStrategy.LATEST)
client.cluster.leader_for_partition = mock.MagicMock()
client.cluster.leader_for_partition.side_effect = [None, -1, 0]
yield from fetcher.update_fetch_positions([partition])
self.assertEqual(state.position, 4)
client.cluster.leader_for_partition = mock.MagicMock()
client.cluster.leader_for_partition.return_value = 1
client.send = mock.MagicMock()
client.send.side_effect = asyncio.coroutine(
lambda n, r: OffsetResponse([('test', [(0, 3, [])])]))
state.await_reset(OffsetResetStrategy.LATEST)
with self.assertRaises(UnknownTopicOrPartitionError):
yield from fetcher.update_fetch_positions([partition])
client.send.side_effect = asyncio.coroutine(
lambda n, r: OffsetResponse([('test', [(0, -1, [])])]))
with self.assertRaises(UnknownError):
yield from fetcher.update_fetch_positions([partition])
yield from fetcher.close()
def test_update_fetch_positions(self):
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
subscriptions = SubscriptionState('latest')
fetcher = Fetcher(client, subscriptions, loop=self.loop)
partition = TopicPartition('test', 0)
# partition is not assigned, should be ignored
yield from fetcher.update_fetch_positions([partition])
state = TopicPartitionState()
state.seek(0)
subscriptions.assignment[partition] = state
# partition is fetchable, no need to update position
yield from fetcher.update_fetch_positions([partition])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
client.send.side_effect = asyncio.coroutine(
lambda n, r: OffsetResponse[0]([('test', [(0, 0, [4])])]))
state.await_reset(OffsetResetStrategy.LATEST)
client.cluster.leader_for_partition = mock.MagicMock()
client.cluster.leader_for_partition.side_effect = [None, -1, 0]
yield from fetcher.update_fetch_positions([partition])
self.assertEqual(state.position, 4)
client.cluster.leader_for_partition = mock.MagicMock()
client.cluster.leader_for_partition.return_value = 1
client.send = mock.MagicMock()
client.send.side_effect = asyncio.coroutine(
lambda n, r: OffsetResponse[0]([('test', [(0, 3, [])])]))
state.await_reset(OffsetResetStrategy.LATEST)
with self.assertRaises(UnknownTopicOrPartitionError):
yield from fetcher.update_fetch_positions([partition])
client.send.side_effect = asyncio.coroutine(
lambda n, r: OffsetResponse[0]([('test', [(0, -1, [])])]))
with self.assertRaises(UnknownError):
yield from fetcher.update_fetch_positions([partition])
yield from fetcher.close()
async def test_fetcher_offsets_for_times(self):
client = AIOKafkaClient(bootstrap_servers=[])
async def ready(conn):
return True
async def _maybe_wait_metadata():
return False
client.ready = mock.MagicMock()
client.ready.side_effect = ready
client._maybe_wait_metadata = mock.MagicMock()
client._maybe_wait_metadata.side_effect = _maybe_wait_metadata
client.cluster.leader_for_partition = mock.MagicMock()
client.cluster.leader_for_partition.return_value = 0
client._api_version = (0, 10, 1)
subscriptions = SubscriptionState()
fetcher = Fetcher(client, subscriptions)
tp0 = TopicPartition("topic", 0)
tp1 = TopicPartition("topic", 1)
# Timeouting will result in KafkaTimeoutError
with mock.patch.object(fetcher, "_proc_offset_requests") as mocked:
mocked.side_effect = asyncio.TimeoutError
with self.assertRaises(KafkaTimeoutError):
await fetcher.get_offsets_by_times({tp0: 0}, 1000)
# Broker returns UnsupportedForMessageFormatError
with mock.patch.object(client, "send") as mocked:
async def mock_send(node_id, request):
return OffsetResponse[1]([("topic", [(0, 43, -1, -1)]),
("topic", [(1, 0, 1000, 9999)])])
mocked.side_effect = mock_send
offsets = await fetcher.get_offsets_by_times({
tp0: 0,
tp1: 0
}, 1000)
self.assertEqual(offsets, {
tp0: None,
tp1: OffsetAndTimestamp(9999, 1000),
})
# Brokers returns NotLeaderForPartitionError
with mock.patch.object(client, "send") as mocked:
async def mock_send(node_id, request):
return OffsetResponse[1]([
("topic", [(0, 6, -1, -1)]),
])
mocked.side_effect = mock_send
with self.assertRaises(NotLeaderForPartitionError):
await fetcher._proc_offset_request(0, {"topic": (0, 1000)})
# Broker returns UnknownTopicOrPartitionError
with mock.patch.object(client, "send") as mocked:
async def mock_send(node_id, request):
return OffsetResponse[1]([
("topic", [(0, 3, -1, -1)]),
])
mocked.side_effect = mock_send
with self.assertLogs("aiokafka.consumer.fetcher", "WARN") as cm:
with self.assertRaises(UnknownTopicOrPartitionError):
await fetcher._proc_offset_request(0, {"topic": (0, 1000)})
if cm is not None:
self.assertIn("Received unknown topic or partition error",
cm.output[0])
async def test_proc_fetch_request(self):
client = AIOKafkaClient(bootstrap_servers=[])
subscriptions = SubscriptionState()
fetcher = Fetcher(client, subscriptions, auto_offset_reset="latest")
tp = TopicPartition('test', 0)
tp_info = (tp.topic, [(tp.partition, 4, 100000)])
req = FetchRequest(
-1, # replica_id
100,
100,
[tp_info])
async def ready(conn):
return True
def force_metadata_update():
fut = create_future()
fut.set_result(False)
return fut
client.ready = mock.MagicMock()
client.ready.side_effect = ready
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = force_metadata_update
client.send = mock.MagicMock()
builder = LegacyRecordBatchBuilder(magic=1,
compression_type=0,
batch_size=99999999)
builder.append(offset=4, value=b"test msg", key=None, timestamp=None)
raw_batch = bytes(builder.build())
fetch_response = FetchResponse([('test', [(0, 0, 9, raw_batch)])])
async def send(node, request):
nonlocal fetch_response
return fetch_response
client.send.side_effect = send
subscriptions.assign_from_user({tp})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(tp)
# The partition has no active position, so will ignore result
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(fetcher._records, {})
# The partition's position does not match request's fetch offset
subscriptions.seek(tp, 0)
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(fetcher._records, {})
subscriptions.seek(tp, 4)
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone().value, b"test msg")
# If position changed after fetch request passed
subscriptions.seek(tp, 4)
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
subscriptions.seek(tp, 10)
self.assertIsNone(buf.getone())
# If assignment is lost after fetch request passed
subscriptions.seek(tp, 4)
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
subscriptions.unsubscribe()
self.assertIsNone(buf.getone())
subscriptions.assign_from_user({tp})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(tp)
# error -> no partition found (UnknownTopicOrPartitionError)
subscriptions.seek(tp, 4)
fetcher._records.clear()
fetch_response = FetchResponse([('test', [(0, 3, 9, raw_batch)])])
cc = client.force_metadata_update.call_count
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(client.force_metadata_update.call_count, cc + 1)
# error -> topic auth failed (TopicAuthorizationFailedError)
fetch_response = FetchResponse([('test', [(0, 29, 9, raw_batch)])])
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(TopicAuthorizationFailedError):
await fetcher.next_record([])
# error -> unknown
fetch_response = FetchResponse([('test', [(0, -1, 9, raw_batch)])])
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, False)
# error -> offset out of range with offset strategy
fetch_response = FetchResponse([('test', [(0, 1, 9, raw_batch)])])
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(tp_state.has_valid_position, False)
self.assertEqual(tp_state.awaiting_reset, True)
self.assertEqual(tp_state.reset_strategy, OffsetResetStrategy.LATEST)
# error -> offset out of range without offset strategy
subscriptions.seek(tp, 4)
fetcher._default_reset_strategy = OffsetResetStrategy.NONE
needs_wake_up = await fetcher._proc_fetch_request(assignment, 0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(OffsetOutOfRangeError):
await fetcher.next_record([])
await fetcher.close()
def test_proc_fetch_request(self):
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
subscriptions = SubscriptionState('latest')
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp = TopicPartition('test', 0)
tp_info = (tp.topic, [(tp.partition, 155, 100000)])
req = FetchRequest(
-1, # replica_id
100, 100, [tp_info])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
msg = Message(b"test msg")
msg._encode_self()
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 0, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
state = TopicPartitionState()
state.seek(0)
subscriptions.assignment[tp] = state
subscriptions.needs_partition_assignment = False
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone(), None) # invalid offset, msg is ignored
state.seek(4)
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone().value, b"test msg")
# error -> no partition found
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 3, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
# error -> topic auth failed
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 29, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(TopicAuthorizationFailedError):
yield from fetcher.next_record([])
# error -> unknown
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, -1, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
# error -> offset out of range
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 1, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(state.is_fetchable(), False)
state.seek(4)
subscriptions._default_offset_reset_strategy = OffsetResetStrategy.NONE
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 1, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(OffsetOutOfRangeError):
yield from fetcher.next_record([])
yield from fetcher.close()
def test_proc_fetch_request(self):
client = AIOKafkaClient(loop=self.loop, bootstrap_servers=[])
subscriptions = SubscriptionState('latest')
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp = TopicPartition('test', 0)
tp_info = (tp.topic, [(tp.partition, 155, 100000)])
req = FetchRequest(
-1, # replica_id
100,
100,
[tp_info])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
msg = Message(b"test msg")
msg._encode_self()
client.send.side_effect = asyncio.coroutine(lambda n, r: FetchResponse(
[('test', [(0, 0, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
state = TopicPartitionState()
state.seek(0)
subscriptions.assignment[tp] = state
subscriptions.needs_partition_assignment = False
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone(), None) # invalid offset, msg is ignored
state.seek(4)
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone().value, b"test msg")
# error -> no partition found
client.send.side_effect = asyncio.coroutine(lambda n, r: FetchResponse(
[('test', [(0, 3, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
# error -> topic auth failed
client.send.side_effect = asyncio.coroutine(lambda n, r: FetchResponse(
[('test', [(0, 29, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(TopicAuthorizationFailedError):
yield from fetcher.next_record([])
# error -> unknown
client.send.side_effect = asyncio.coroutine(lambda n, r: FetchResponse(
[('test', [(0, -1, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
# error -> offset out of range
client.send.side_effect = asyncio.coroutine(lambda n, r: FetchResponse(
[('test', [(0, 1, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(state.is_fetchable(), False)
state.seek(4)
subscriptions._default_offset_reset_strategy = OffsetResetStrategy.NONE
client.send.side_effect = asyncio.coroutine(lambda n, r: FetchResponse(
[('test', [(0, 1, 9, [(4, 10, msg)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(OffsetOutOfRangeError):
yield from fetcher.next_record([])
yield from fetcher.close()
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_compacted_topic_consumption(self):
# Compacted topics can have offsets skipped
client = AIOKafkaClient(loop=self.loop, bootstrap_servers=[])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
subscriptions = SubscriptionState('latest')
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp = TopicPartition('test', 0)
req = FetchRequest(
-1, # replica_id
100,
100,
[(tp.topic, [(tp.partition, 155, 100000)])])
msg1 = Message(b"12345", key=b"1")
msg1._encode_self()
msg2 = Message(b"23456", key=b"2")
msg2._encode_self()
msg3 = Message(b"34567", key=b"3")
msg3._encode_self()
resp = FetchResponse([(
'test',
[(
0,
0,
3000, # partition, error_code, highwater_offset
[
(160, 5, msg1), # offset, len_bytes, bytes
(162, 5, msg2),
(167, 5, msg3),
])])])
client.send.side_effect = asyncio.coroutine(lambda n, r: resp)
state = TopicPartitionState()
state.seek(155)
state.drop_pending_message_set = False
subscriptions.assignment[tp] = state
subscriptions.needs_partition_assignment = False
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
# Test successful getone
first = buf.getone()
self.assertEqual(state.position, 161)
self.assertEqual((first.value, first.key, first.offset),
(msg1.value, msg1.key, 160))
# Test successful getmany
second, third = buf.getall()
self.assertEqual(state.position, 168)
self.assertEqual((second.value, second.key, second.offset),
(msg2.value, msg2.key, 162))
self.assertEqual((third.value, third.key, third.offset),
(msg3.value, msg3.key, 167))
def test_proc_fetch_request(self):
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
subscriptions = SubscriptionState('latest')
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp = TopicPartition('test', 0)
tp_info = (tp.topic, [(tp.partition, 155, 100000)])
req = FetchRequest(
-1, # replica_id
100, 100, [tp_info])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
builder = LegacyRecordBatchBuilder(
magic=1, compression_type=0, batch_size=99999999)
builder.append(offset=4, value=b"test msg", key=None, timestamp=None)
raw_batch = bytes(builder.build())
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 0, 9, raw_batch)])]))
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
state = TopicPartitionState()
state.seek(0)
subscriptions.assignment[tp] = state
subscriptions.needs_partition_assignment = False
fetcher._in_flight.add(0)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone(), None) # invalid offset, msg is ignored
state.seek(4)
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone().value, b"test msg")
# error -> no partition found
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 3, 9, raw_batch)])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
# error -> topic auth failed
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 29, 9, raw_batch)])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(TopicAuthorizationFailedError):
yield from fetcher.next_record([])
# error -> unknown
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, -1, 9, raw_batch)])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
# error -> offset out of range with offset strategy
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 1, 9, raw_batch)])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
with mock.patch.object(fetcher, "update_fetch_positions") as mocked:
mocked.side_effect = asyncio.coroutine(lambda o: None)
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(state.is_fetchable(), False)
mocked.assert_called_with([tp])
# error -> offset out of range with strategy errors out
state.seek(4)
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 1, 9, [(4, 10, raw_batch)])])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
with mock.patch.object(fetcher, "update_fetch_positions") as mocked:
# the exception should not fail execution here
@asyncio.coroutine
def mock_async_raises(offests):
raise Exception()
mocked.side_effect = mock_async_raises
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(state.is_fetchable(), False)
mocked.assert_called_with([tp])
# error -> offset out of range without offset strategy
state.seek(4)
subscriptions._default_offset_reset_strategy = OffsetResetStrategy.NONE
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 1, 9, raw_batch)])]))
fetcher._in_flight.add(0)
fetcher._records.clear()
needs_wake_up = yield from fetcher._proc_fetch_request(0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(OffsetOutOfRangeError):
yield from fetcher.next_record([])
yield from fetcher.close()
def test_fetcher_offsets_for_times(self):
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client._maybe_wait_metadata = mock.MagicMock()
client._maybe_wait_metadata.side_effect = asyncio.coroutine(
lambda: False)
client.cluster.leader_for_partition = mock.MagicMock()
client.cluster.leader_for_partition.return_value = 0
client._api_version = (0, 10, 1)
subscriptions = SubscriptionState(loop=self.loop)
fetcher = Fetcher(client, subscriptions, loop=self.loop)
tp0 = TopicPartition("topic", 0)
tp1 = TopicPartition("topic", 1)
# Timeouting will result in KafkaTimeoutError
with mock.patch.object(fetcher, "_proc_offset_requests") as mocked:
mocked.side_effect = asyncio.TimeoutError
with self.assertRaises(KafkaTimeoutError):
yield from fetcher.get_offsets_by_times({tp0: 0}, 1000)
# Broker returns UnsupportedForMessageFormatError
with mock.patch.object(client, "send") as mocked:
@asyncio.coroutine
def mock_send(node_id, request):
return OffsetResponse[1]([
("topic", [(0, 43, -1, -1)]),
("topic", [(1, 0, 1000, 9999)])
])
mocked.side_effect = mock_send
offsets = yield from fetcher.get_offsets_by_times(
{tp0: 0, tp1: 0}, 1000)
self.assertEqual(offsets, {
tp0: None,
tp1: OffsetAndTimestamp(9999, 1000),
})
# Brokers returns NotLeaderForPartitionError
with mock.patch.object(client, "send") as mocked:
@asyncio.coroutine
def mock_send(node_id, request):
return OffsetResponse[1]([
("topic", [(0, 6, -1, -1)]),
])
mocked.side_effect = mock_send
with self.assertRaises(NotLeaderForPartitionError):
yield from fetcher._proc_offset_request(
0, {"topic": (0, 1000)})
# Broker returns UnknownTopicOrPartitionError
with mock.patch.object(client, "send") as mocked:
@asyncio.coroutine
def mock_send(node_id, request):
return OffsetResponse[1]([
("topic", [(0, 3, -1, -1)]),
])
mocked.side_effect = mock_send
with self.assertLogs("aiokafka.consumer.fetcher", "WARN") as cm:
with self.assertRaises(UnknownTopicOrPartitionError):
yield from fetcher._proc_offset_request(
0, {"topic": (0, 1000)})
if cm is not None:
self.assertIn(
"Received unknown topic or partition error", cm.output[0])
def test_proc_fetch_request(self):
client = AIOKafkaClient(
loop=self.loop,
bootstrap_servers=[])
subscriptions = SubscriptionState(loop=self.loop)
fetcher = Fetcher(
client, subscriptions, auto_offset_reset="latest", loop=self.loop)
tp = TopicPartition('test', 0)
tp_info = (tp.topic, [(tp.partition, 4, 100000)])
req = FetchRequest(
-1, # replica_id
100, 100, [tp_info])
client.ready = mock.MagicMock()
client.ready.side_effect = asyncio.coroutine(lambda a: True)
client.force_metadata_update = mock.MagicMock()
client.force_metadata_update.side_effect = asyncio.coroutine(
lambda: False)
client.send = mock.MagicMock()
builder = LegacyRecordBatchBuilder(
magic=1, compression_type=0, batch_size=99999999)
builder.append(offset=4, value=b"test msg", key=None, timestamp=None)
raw_batch = bytes(builder.build())
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 0, 9, raw_batch)])]))
subscriptions.assign_from_user({tp})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(tp)
# The partition has no active position, so will ignore result
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(fetcher._records, {})
# The partition's position does not match request's fetch offset
subscriptions.seek(tp, 0)
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(fetcher._records, {})
subscriptions.seek(tp, 4)
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, True)
buf = fetcher._records[tp]
self.assertEqual(buf.getone().value, b"test msg")
# If position changed after fetch request passed
subscriptions.seek(tp, 4)
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
subscriptions.seek(tp, 10)
self.assertIsNone(buf.getone())
# If assignment is lost after fetch request passed
subscriptions.seek(tp, 4)
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
subscriptions.unsubscribe()
self.assertIsNone(buf.getone())
subscriptions.assign_from_user({tp})
assignment = subscriptions.subscription.assignment
tp_state = assignment.state_value(tp)
# error -> no partition found (UnknownTopicOrPartitionError)
subscriptions.seek(tp, 4)
fetcher._records.clear()
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 3, 9, raw_batch)])]))
cc = client.force_metadata_update.call_count
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(client.force_metadata_update.call_count, cc + 1)
# error -> topic auth failed (TopicAuthorizationFailedError)
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 29, 9, raw_batch)])]))
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(TopicAuthorizationFailedError):
yield from fetcher.next_record([])
# error -> unknown
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, -1, 9, raw_batch)])]))
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, False)
# error -> offset out of range with offset strategy
client.send.side_effect = asyncio.coroutine(
lambda n, r: FetchResponse(
[('test', [(0, 1, 9, raw_batch)])]))
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, False)
self.assertEqual(tp_state.has_valid_position, False)
self.assertEqual(tp_state.awaiting_reset, True)
self.assertEqual(tp_state.reset_strategy, OffsetResetStrategy.LATEST)
# error -> offset out of range without offset strategy
subscriptions.seek(tp, 4)
fetcher._default_reset_strategy = OffsetResetStrategy.NONE
needs_wake_up = yield from fetcher._proc_fetch_request(
assignment, 0, req)
self.assertEqual(needs_wake_up, True)
with self.assertRaises(OffsetOutOfRangeError):
yield from fetcher.next_record([])
yield from fetcher.close()