Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
await consumer.start()
self.add_cleanup(consumer.stop)
await consumer.seek_to_committed()
result = []
for i in range(20):
msg = await consumer.getone()
result.append(msg.value)
self.assertEqual(set(available_msgs), set(result))
await consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
await consumer.seek_to_committed(TopicPartition(self.topic, 0))
msg = await consumer.getone(TopicPartition(self.topic, 0))
self.assertEqual(msg.value, b'9')
await consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
await consumer.stop()
# subscribe by topic
consumer = AIOKafkaConsumer(
loop=self.loop, group_id='test-group',
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=(self.topic,))
await consumer.start()
self.add_cleanup(consumer.stop)
await consumer.seek_to_committed()
result = []
for i in range(10):
msg = await consumer.getone()
result.append(msg.value)
self.assertEqual(set(msgs2), set(result))
consumer = AIOKafkaConsumer(
loop=self.loop, group_id='test-group',
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(pattern="topic-test_manual_subs*")
await consumer.start()
self.add_cleanup(consumer.stop)
await consumer.seek_to_committed()
result = []
for i in range(20):
msg = await consumer.getone()
result.append(msg.value)
self.assertEqual(set(available_msgs), set(result))
await consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
await consumer.seek_to_committed(TopicPartition(self.topic, 0))
msg = await consumer.getone(TopicPartition(self.topic, 0))
self.assertEqual(msg.value, b'9')
await consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
await consumer.stop()
# subscribe by topic
consumer = AIOKafkaConsumer(
loop=self.loop, group_id='test-group',
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=(self.topic,))
await consumer.start()
self.add_cleanup(consumer.stop)
await consumer.seek_to_committed()
async def test_consumer_commit_validation(self):
consumer = await self.consumer_factory()
self.add_cleanup(consumer.stop)
tp = TopicPartition(self.topic, 0)
offset = await consumer.position(tp)
offset_and_metadata = OffsetAndMetadata(offset, "")
with self.assertRaises(ValueError):
await consumer.commit({})
with self.assertRaises(ValueError):
await consumer.commit("something")
with self.assertRaises(ValueError):
await consumer.commit({tp: (offset, "metadata", 100)})
with self.assertRaisesRegex(
ValueError, "Key should be TopicPartition instance"):
await consumer.commit({"my_topic": offset_and_metadata})
with self.assertRaisesRegex(
ValueError, "Metadata should be a string"):
await consumer.commit({tp: (offset, 1000)})
with self.assertRaisesRegex(
ValueError, "Metadata should be a string"):
await consumer.commit({tp: (offset, b"\x00\x02")})
async def test_commit_failed_scenarios(self):
client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
await client.bootstrap()
await self.wait_topic(client, 'topic1')
subscription = SubscriptionState(loop=self.loop)
subscription.subscribe(topics=set(['topic1']))
coordinator = GroupCoordinator(
client, subscription, loop=self.loop,
group_id='test-offsets-group')
await subscription.wait_for_assignment()
assignment = subscription.subscription.assignment
offsets = {TopicPartition('topic1', 0): OffsetAndMetadata(1, '')}
await coordinator.commit_offsets(assignment, offsets)
_orig_send_req = coordinator._send_req
with mock.patch.object(coordinator, "_send_req") as mocked:
commit_error = None
async def mock_send_req(request):
if request.API_KEY == OffsetCommitRequest[0].API_KEY:
if isinstance(commit_error, list):
error_code = commit_error.pop(0).errno
else:
error_code = commit_error.errno
resp_topics = [("topic1", [(0, error_code)])]
return OffsetCommitResponse_v2(resp_topics)
return (await _orig_send_req(request))
mocked.side_effect = mock_send_req
# Success case
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, True)
self.assertEqual(fut.result(), OffsetAndMetadata(12, ""))
# Calling again will fast return without a request
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, True)
self.assertEqual(mocked.call_count, 1)
# Commit not found case
fetched_offsets = {}
assignment, tp_state, fut = reset_assignment()
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, True)
self.assertEqual(fut.result(), OffsetAndMetadata(-1, ""))
# Retriable error will be skipped
assignment, tp_state, fut = reset_assignment()
mocked.side_effect = Errors.GroupCoordinatorNotAvailableError()
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, False)
# Not retriable error will not be skipped
mocked.side_effect = Errors.UnknownError()
with self.assertRaises(Errors.UnknownError):
await coordinator._maybe_refresh_commit_offsets(assignment)
test_self.assertEqual(need_update, [tp])
return fetched_offsets
mocked.side_effect = do_fetch
def reset_assignment():
subscription.assign_from_user({tp})
assignment = subscription.subscription.assignment
tp_state = assignment.state_value(tp)
fut = tp_state.fetch_committed()
return assignment, tp_state, fut
assignment, tp_state, fut = reset_assignment()
# Success case
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, True)
self.assertEqual(fut.result(), OffsetAndMetadata(12, ""))
# Calling again will fast return without a request
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, True)
self.assertEqual(mocked.call_count, 1)
# Commit not found case
fetched_offsets = {}
assignment, tp_state, fut = reset_assignment()
resp = await coordinator._maybe_refresh_commit_offsets(assignment)
self.assertEqual(resp, True)
self.assertEqual(fut.result(), OffsetAndMetadata(-1, ""))
# Retriable error will be skipped
assignment, tp_state, fut = reset_assignment()
mocked.side_effect = Errors.GroupCoordinatorNotAvailableError()
mocked.side_effect = mock_send_req
# 0 partitions call should just fast return
res = await coordinator.fetch_committed_offsets({})
self.assertEqual(res, {})
self.assertEqual(mocked.call_count, 0)
fetch_error = [
Errors.GroupLoadInProgressError,
Errors.GroupLoadInProgressError,
Errors.NoError,
Errors.NoError,
Errors.NoError
]
res = await coordinator.fetch_committed_offsets(partitions)
self.assertEqual(res, {tp: OffsetAndMetadata(10, "")})
# Just omit the topic with a warning
fetch_error = Errors.UnknownTopicOrPartitionError
res = await coordinator.fetch_committed_offsets(partitions)
self.assertEqual(res, {})
fetch_error = [
Errors.NotCoordinatorForGroupError,
Errors.NotCoordinatorForGroupError,
Errors.NoError,
Errors.NoError,
Errors.NoError
]
r = await coordinator.fetch_committed_offsets(partitions)
self.assertEqual(r, {tp: OffsetAndMetadata(10, "")})
self.assertIsNone(tp_state._position)
await fetcher._update_fetch_positions(assignment, 0, [partition])
self.assertEqual(tp_state._position, 12)
self.assertEqual(fetcher._proc_offset_request.call_count, 1)
# CASE: seeked while waiting for committed to be resolved
assignment, tp_state = reset_assignment()
update_task = ensure_future(
fetcher._update_fetch_positions(assignment, 0, [partition]),
loop=self.loop
)
await asyncio.sleep(0.1, loop=self.loop)
self.assertFalse(update_task.done())
tp_state.seek(8)
tp_state.update_committed(OffsetAndMetadata(4, ""))
await update_task
self.assertEqual(tp_state._position, 8)
self.assertEqual(fetcher._proc_offset_request.call_count, 1)
# CASE: awaiting_reset during waiting for committed
assignment, tp_state = reset_assignment()
update_task = ensure_future(
fetcher._update_fetch_positions(assignment, 0, [partition]),
loop=self.loop
)
await asyncio.sleep(0.1, loop=self.loop)
self.assertFalse(update_task.done())
tp_state.await_reset(OffsetResetStrategy.LATEST)
tp_state.update_committed(OffsetAndMetadata(4, ""))
await update_task
async def test_sender__do_txn_offset_commit_ok(self):
sender = await self._setup_sender()
offsets = {
TopicPartition("topic", 0): OffsetAndMetadata(10, ""),
TopicPartition("topic", 1): OffsetAndMetadata(11, ""),
}
add_handler = TxnOffsetCommitHandler(sender, offsets, "some_group")
tm = sender._txn_manager
tm.offset_committed = mock.Mock()
# Handle response
cls = TxnOffsetCommitResponse[0]
resp = cls(
throttle_time_ms=300,
errors=[
("topic", [
(0, NoError.errno),
(1, NoError.errno)
])
]
def all_consumed_offsets(self) -> Dict[TopicPartition, OffsetAndMetadata]:
""" Returns consumed offsets as {TopicPartition: OffsetAndMetadata} """
all_consumed = {}
for tp in self._topic_partitions:
state = self.state_value(tp)
if state.has_valid_position:
all_consumed[tp] = OffsetAndMetadata(state.position, '')
return all_consumed