def test_process_consumer_offset_message_invalid_message(self, parse_mock): parse_mock.side_effect = InvalidMessageException kafka_group_reader = KafkaGroupReader(mock.Mock()) message = mock.MagicMock(spec=ConsumerRecord) kafka_group_reader.process_consumer_offset_message(message) assert kafka_group_reader._kafka_groups == dict()
def test_remove_unsubscribed_topics(self): kafka_config = mock.Mock() kafka_group_reader = KafkaGroupReader(kafka_config) kafka_group_reader._kafka_groups = { 'my_group1': { 'topic1': {0: 0, 1: 1}, 'topic2': {0: 1}, }, 'my_group2': { 'topic1': {0: 0, 1: 0}, 'topic2': {0: 1}, }, } expected = { 'my_group1': { 'topic1': {0: 0, 1: 1}, 'topic2': {0: 1}, }, 'my_group2': { 'topic2': {0: 1}, }, } kafka_group_reader._remove_unsubscribed_topics() assert kafka_group_reader._kafka_groups == expected
def test_parse_consumer_offset_message_no_value(self): kafka_config = mock.Mock() kafka_group_reader = KafkaGroupReader(kafka_config) message = Message(0, '__consumer_offsets', self.key_ok, None) group, topic, partition, offset = kafka_group_reader.parse_consumer_offset_message(message) assert group == 'group1' assert topic == 'topic1' assert partition == 15 assert offset is None
def get_topics_for_group_from_kafka( cls, cluster_config, groupid, ): kafka_group_reader = KafkaGroupReader(cluster_config) group_topics = kafka_group_reader.read_group(groupid) if len(group_topics) == 0: print( "Error: Either Consumer Group {groupid} has never committed offsets, or the group does not exist.".format( groupid=groupid, ), file=sys.stderr, ) return group_topics
def test_process_consumer_offset_message_topic_get(self): kafka_config = mock.Mock() kafka_group_reader = KafkaGroupReader(kafka_config) assert kafka_group_reader._kafka_groups == {} with mock.patch.object( kafka_group_reader, 'parse_consumer_offset_message', return_value=[ 'test_group', 'test_topic', 0, 45 ], autospec=True ): kafka_group_reader.process_consumer_offset_message('test message') assert kafka_group_reader._kafka_groups['test_group'] == {'test_topic': {0: 45}}
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_process_consumer_offset_message_group(self, parse_mock): parse_mock.side_effect = [('test.a', 'topic1', 0, 123), ('test.a', 'topic1', 1, 124), ('test.a', 'topic2', 0, 125), ('my_test', 'topic1', 0, 123), ('my_test', 'topic2', 0, 124), ('my_test', 'topic2', 0, None), ('my_test2', 'topic3', 0, 123), ] kafka_group_reader = KafkaGroupReader(mock.Mock()) for _ in range(7): message = mock.MagicMock(spec=ConsumerRecord) kafka_group_reader.process_consumer_offset_message(message) expected = { 'test.a': { 'topic1': {0: 123, 1: 124}, 'topic2': {0: 125}, }, 'my_test2': { 'topic3': {0: 123}, }, 'my_test': { 'topic1': {0: 123}, }, } # Convert the defaultdict to a normal dict for comparison actual = {} for group, topics in six.iteritems(kafka_group_reader._kafka_groups): actual[group] = {} for topic, partitions in six.iteritems(topics): actual[group][topic] = {} for partition, offset in six.iteritems(partitions): actual[group][topic][partition] = offset assert actual == expected
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 get_kafka_groups(cls, cluster_config): '''Get the group_id of groups committed into Kafka.''' kafka_group_reader = KafkaGroupReader(cluster_config) return list(kafka_group_reader.read_groups().keys())
def get_topics_for_group_from_kafka(cls, cluster_config, groupid): kafka_group_reader = KafkaGroupReader(cluster_config) return kafka_group_reader.read_group(groupid)
def get_kafka_groups(cls, cluster_config): '''Get the group_id of groups committed into Kafka.''' kafka_group_reader = KafkaGroupReader(cluster_config) return kafka_group_reader.read_groups().keys()