import unittest from datetime import datetime from kafka import * # noqa from kafka.common import * # noqa from kafka.consumer import MAX_FETCH_BUFFER_SIZE_BYTES from fixtures import ZookeeperFixture, KafkaFixture from testutil import * @unittest.skipIf(skip_integration(), 'Skipping Integration') class TestConsumerIntegration(KafkaIntegrationTestCase): @classmethod def setUpClass(cls): cls.zk = ZookeeperFixture.instance() cls.server1 = KafkaFixture.instance(0, cls.zk.host, cls.zk.port) cls.server2 = KafkaFixture.instance(1, cls.zk.host, cls.zk.port) cls.server = cls.server1 # Bootstrapping server @classmethod def tearDownClass(cls): # noqa cls.server1.close() cls.server2.close() cls.zk.close() def send_messages(self, partition, messages): messages = [ create_message(self.msg(str(msg))) for msg in messages ] produce = ProduceRequest(self.topic, partition, messages = messages) resp, = self.client.send_produce_request([produce]) self.assertEquals(resp.error, 0) return [ x.value for x in messages ] def assert_message_count(self, messages, num_messages): # Make sure we got them all self.assertEquals(len(messages), num_messages) # Make sure there are no duplicates self.assertEquals(len(set(messages)), num_messages) def test_simple_consumer(self): self.send_messages(0, range(0, 100)) self.send_messages(1, range(100, 200)) # Start a consumer consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=False, iter_timeout=0) self.assert_message_count([ message for message in consumer ], 200) consumer.stop() def test_simple_consumer__seek(self): self.send_messages(0, range(0, 100)) self.send_messages(1, range(100, 200)) consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=False, iter_timeout=0) # Rewind 10 messages from the end consumer.seek(-10, 2) self.assert_message_count([ message for message in consumer ], 10) # Rewind 13 messages from the end consumer.seek(-13, 2) self.assert_message_count([ message for message in consumer ], 13) consumer.stop() def test_simple_consumer_blocking(self): consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=False, iter_timeout=0) # Ask for 5 messages, nothing in queue, block 5 seconds with Timer() as t: messages = consumer.get_messages(block=True, timeout=5) self.assert_message_count(messages, 0) self.assertGreaterEqual(t.interval, 5) self.send_messages(0, range(0, 10)) # Ask for 5 messages, 10 in queue. Get 5 back, no blocking with Timer() as t: messages = consumer.get_messages(count=5, block=True, timeout=5) self.assert_message_count(messages, 5) self.assertLessEqual(t.interval, 1) # Ask for 10 messages, get 5 back, block 5 seconds with Timer() as t: messages = consumer.get_messages(count=10, block=True, timeout=5) self.assert_message_count(messages, 5) self.assertGreaterEqual(t.interval, 5) consumer.stop() def test_simple_consumer_pending(self): # Produce 10 messages to partitions 0 and 1 self.send_messages(0, range(0, 10)) self.send_messages(1, range(10, 20)) consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=False, iter_timeout=0) self.assertEquals(consumer.pending(), 20) self.assertEquals(consumer.pending(partitions=[0]), 10) self.assertEquals(consumer.pending(partitions=[1]), 10) consumer.stop() def test_multi_process_consumer(self): # Produce 100 messages to partitions 0 and 1 self.send_messages(0, range(0, 100)) self.send_messages(1, range(100, 200)) consumer = MultiProcessConsumer(self.client, "grp1", self.topic, auto_commit=False) self.assert_message_count([ message for message in consumer ], 200) consumer.stop() def test_multi_process_consumer_blocking(self): consumer = MultiProcessConsumer(self.client, "grp1", self.topic, auto_commit=False) # Ask for 5 messages, No messages in queue, block 5 seconds with Timer() as t: messages = consumer.get_messages(block=True, timeout=5) self.assert_message_count(messages, 0) self.assertGreaterEqual(t.interval, 5) # Send 10 messages self.send_messages(0, range(0, 10)) # Ask for 5 messages, 10 messages in queue, block 0 seconds with Timer() as t: messages = consumer.get_messages(count=5, block=True, timeout=5) self.assert_message_count(messages, 5) self.assertLessEqual(t.interval, 1) # Ask for 10 messages, 5 in queue, block 5 seconds with Timer() as t: messages = consumer.get_messages(count=10, block=True, timeout=5) self.assert_message_count(messages, 5) self.assertGreaterEqual(t.interval, 5) consumer.stop() def test_multi_proc_pending(self): self.send_messages(0, range(0, 10)) self.send_messages(1, range(10, 20)) consumer = MultiProcessConsumer(self.client, "group1", self.topic, auto_commit=False) self.assertEquals(consumer.pending(), 20) self.assertEquals(consumer.pending(partitions=[0]), 10) self.assertEquals(consumer.pending(partitions=[1]), 10) consumer.stop() def test_large_messages(self): # Produce 10 "normal" size messages small_messages = self.send_messages(0, [ str(x) for x in range(10) ]) # Produce 10 messages that are large (bigger than default fetch size) large_messages = self.send_messages(0, [ random_string(5000) for x in range(10) ]) # Consumer should still get all of them consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=False, iter_timeout=0) expected_messages = set(small_messages + large_messages) actual_messages = set([ x.message.value for x in consumer ]) self.assertEqual(expected_messages, actual_messages) consumer.stop() def test_huge_messages(self): huge_message, = self.send_messages(0, [ create_message(random_string(MAX_FETCH_BUFFER_SIZE_BYTES + 10)), ]) # Create a consumer with the default buffer size consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=False, iter_timeout=0) # This consumer failes to get the message with self.assertRaises(ConsumerFetchSizeTooSmall): consumer.get_message(False, 0.1) consumer.stop() # Create a consumer with no fetch size limit big_consumer = SimpleConsumer(self.client, "group1", self.topic, max_buffer_size=None, partitions=[0], auto_commit=False, iter_timeout=0) # Seek to the last message big_consumer.seek(-1, 2) # Consume giant message successfully message = big_consumer.get_message(block=False, timeout=10) self.assertIsNotNone(message) self.assertEquals(message.message.value, huge_message) big_consumer.stop() @kafka_versions("0.8.1") def test_offset_behavior__resuming_behavior(self): msgs1 = self.send_messages(0, range(0, 100)) msgs2 = self.send_messages(1, range(100, 200)) # Start a consumer consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=True, auto_commit_every_n=20, iter_timeout=0) # Grab the first 195 messages output_msgs1 = [ consumer.get_message().message.value for _ in xrange(195) ] self.assert_message_count(output_msgs1, 195) consumer.stop() # The offset should be at 180 consumer = SimpleConsumer(self.client, "group1", self.topic, auto_commit=True, auto_commit_every_n=20, iter_timeout=0) # 180-200 self.assert_message_count([ message for message in consumer ], 20) consumer.stop()