def test_get_current_consumer_offsets_invalid_partition_subset(self, kafka_client_mock):
with pytest.raises(UnknownPartitions):
get_current_consumer_offsets(
kafka_client_mock,
self.group,
{'topic1': [1, 99]},
)
def test_get_current_consumer_offsets_unknown_topic(self, kafka_client_mock):
with pytest.raises(UnknownTopic):
get_current_consumer_offsets(
kafka_client_mock,
"this won't even be consulted",
["something that doesn't exist"],
)
def test_get_current_consumer_offsets_invalid_arguments(self, kafka_client_mock):
with pytest.raises(TypeError):
get_current_consumer_offsets(
kafka_client_mock,
"this won't even be consulted",
"this should be a list or dict",
)
def test_get_current_consumer_offsets_from_kafka(
self,
topics,
kafka_client_mock
):
kafka_client_mock = mock.Mock(wraps=kafka_client_mock)
get_current_consumer_offsets(
kafka_client_mock,
self.group,
topics,
offset_storage='kafka',
)
assert kafka_client_mock.send_offset_fetch_request.call_count == 0
assert kafka_client_mock.send_offset_fetch_request_kafka.call_count == 1
def test_get_current_consumer_offsets_invalid_storage(
self,
topics,
kafka_client_mock
):
kafka_client_mock = mock.Mock(wraps=kafka_client_mock)
with pytest.raises(InvalidOffsetStorageError):
get_current_consumer_offsets(
kafka_client_mock,
self.group,
topics,
offset_storage='random_string',
)
assert kafka_client_mock.send_offset_fetch_request.call_count == 0
assert kafka_client_mock.send_offset_fetch_request_kafka.call_count == 0
def test_get_current_consumer_offsets(self, topics, kafka_client_mock):
actual = get_current_consumer_offsets(
kafka_client_mock,
self.group,
topics
)
assert actual == {'topic1': {0: 30, 1: 20, 2: 10}}
def test_get_current_consumer_offsets_unknown_partitions_no_fail(self, kafka_client_mock):
actual = get_current_consumer_offsets(
kafka_client_mock,
self.group,
{'topic1': [99]},
raise_on_error=False
)
assert not actual
def test_get_current_consumer_offsets_unknown_topic_no_fail(self, kafka_client_mock):
actual = get_current_consumer_offsets(
kafka_client_mock,
"this won't even be consulted",
["something that doesn't exist"],
raise_on_error=False
)
assert not actual
def test_get_current_consumer_offsets_invalid_partition_subset_no_fail(self, kafka_client_mock):
actual = get_current_consumer_offsets(
kafka_client_mock,
self.group,
{'topic1': [1, 99]},
raise_on_error=False
)
assert actual['topic1'][1] == 20
# Partition 99 does not exist so it shouldn't be in the result
assert 99 not in actual['topic1']
def get_consumer_offsets_metadata(
kafka_client,
group,
topics,
raise_on_error=True,
offset_storage='zookeeper',
):
"""This method:
* refreshes metadata for the kafka client
* fetches group offsets
* fetches watermarks
:param kafka_client: KafkaClient instance
:param group: group id
:param topics: list of topics
:param raise_on_error: if False the method ignores missing topics and
missing partitions. It still may fail on the request send.
:param offset_storage: String, one of {zookeeper, kafka}.
:returns: dict <topic>: [ConsumerPartitionOffsets]
"""
# Refresh client metadata. We do now use the topic list, because we
# don't want to accidentally create the topic if it does not exist.
# If Kafka is unavailable, let's retry loading client metadata (YELPKAFKA-30)
try:
kafka_client.load_metadata_for_topics()
except KafkaUnavailableError:
kafka_client.load_metadata_for_topics()
group_offsets = get_current_consumer_offsets(kafka_client, group, topics,
raise_on_error,
offset_storage)
watermarks = get_topics_watermarks(kafka_client, topics, raise_on_error)
result = {}
for topic, partitions in six.iteritems(group_offsets):
result[topic] = [
ConsumerPartitionOffsets(
topic=topic,
partition=partition,
current=group_offsets[topic][partition],
highmark=watermarks[topic][partition].highmark,
lowmark=watermarks[topic][partition].lowmark,
) for partition in partitions
]
return result