def test_get_topics_watermarks(self, topics, kafka_client_mock):
actual = get_topics_watermarks(
kafka_client_mock,
topics,
)
assert actual == {'topic1': {
0: PartitionOffsets('topic1', 0, 30, 10),
1: PartitionOffsets('topic1', 1, 30, 5),
2: PartitionOffsets('topic1', 2, 30, 3),
}}
def test_read_groups(self):
kafka_config = mock.Mock()
kafka_group_reader = KafkaGroupReader(kafka_config)
with mock.patch(
'kafka_utils.kafka_consumer_manager.util.KafkaConsumer',
) as mock_consumer:
with mock.patch.object(
kafka_group_reader,
'get_current_watermarks',
return_value={
0: PartitionOffsets('__consumer_offsets', 0, 45, 0),
1: PartitionOffsets('__consumer_offsets', 1, 20, 0),
2: PartitionOffsets('__consumer_offsets', 2, 25, 25),
3: PartitionOffsets('__consumer_offsets', 3, 0, 0),
},
autospec=True):
with mock.patch.object(
kafka_group_reader,
'parse_consumer_offset_message',
side_effect=iter([
('test_group', 'test_topic', 0, 45),
('test_group2', 'test_topic2', 0, 20),
]),
autospec=True,
):
mock_consumer.return_value.__iter__.return_value = iter([
mock.Mock(offset=44, partition=0, topic='test_topic'),
mock.Mock(offset=19, partition=1, topic='test_topic'),
])
mock_consumer.return_value.partitions_for_topic.return_value = [
0, 1
]
kafka_group_reader.read_groups()
assert kafka_group_reader._kafka_groups['test_group'] == {
"test_topic": {
0: 45
}
}
assert kafka_group_reader._kafka_groups['test_group2'] == {
"test_topic2": {
0: 20
}
}
mock_consumer.return_value.assign.call_args_list == [
mock.call([
TopicPartition("__consumer_offsets", 0),
TopicPartition("__consumer_offsets", 1),
]),
mock.call([TopicPartition("__consumer_offsets", 0)]),
]
def test_read_groups_with_partition(self):
kafka_config = mock.Mock()
kafka_group_reader = KafkaGroupReader(kafka_config)
with mock.patch(
'kafka_utils.kafka_consumer_manager.util.KafkaConsumer',
) as mock_consumer:
with mock.patch.object(
kafka_group_reader,
'get_current_watermarks',
return_value={0: PartitionOffsets('test_topic', 0, 45, 0)},
autospec=True):
with mock.patch.object(
kafka_group_reader,
'parse_consumer_offset_message',
return_value=['test_group', 'test_topic', 0, 45],
autospec=True):
mock_consumer.return_value.__iter__.return_value = iter([
mock.Mock(partition=0, topic='test_topic', offset=45)
])
kafka_group_reader.read_groups(partition=0)
assert kafka_group_reader._kafka_groups['test_group'] == {
"test_topic": {
0: 45
}
}
mock_consumer.return_value.assign.assert_called_once_with(
[TopicPartition("__consumer_offsets", 0)])
def test_get_topics_watermarks_invalid_partition_subset_no_fail(self, kafka_client_mock):
actual = get_topics_watermarks(
kafka_client_mock,
{'topic1': [1, 99]},
raise_on_error=False,
)
assert actual['topic1'][1] == PartitionOffsets('topic1', 1, 30, 5)
assert 99 not in actual['topic1']
def test_get_topics_watermarks_commit_error(self, topics, kafka_client_mock):
kafka_client_mock.set_offset_request_error()
actual = get_topics_watermarks(
kafka_client_mock,
{'topic1': [0]},
)
assert actual == {'topic1': {
0: PartitionOffsets('topic1', 0, -1, -1),
}}