def test_manual_subscribe_pattern(self): msgs1 = yield from self.send_messages(0, range(0, 10)) msgs2 = yield from self.send_messages(1, range(10, 20)) available_msgs = msgs1 + msgs2 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*") yield from consumer.start() yield from consumer.seek_to_committed() result = [] for i in range(20): msg = yield from consumer.getone() result.append(msg.value) self.assertEqual(set(available_msgs), set(result)) yield from consumer.commit( {TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')}) yield from consumer.seek_to_committed(TopicPartition(self.topic, 0)) msg = yield from consumer.getone(TopicPartition(self.topic, 0)) self.assertEqual(msg.value, b'9') yield from consumer.commit( {TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')}) yield from 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, )) yield from consumer.start() yield from consumer.seek_to_committed() result = [] for i in range(10): msg = yield from consumer.getone() result.append(msg.value) self.assertEqual(set(msgs2), set(result)) self.assertEqual(consumer.subscription(), set([self.topic])) yield from consumer.stop()
def test_unknown_topic_or_partition(self): consumer = AIOKafkaConsumer( loop=self.loop, group_id=None, bootstrap_servers=self.hosts, auto_offset_reset='earliest', enable_auto_commit=False) yield from consumer.start() with self.assertRaises(UnknownTopicOrPartitionError): yield from consumer.assign([TopicPartition(self.topic, 2222)]) yield from consumer.stop()
def test_consumer_group_without_subscription(self): consumer = AIOKafkaConsumer(loop=self.loop, group_id='group-{}'.format(self.id()), bootstrap_servers=self.hosts, enable_auto_commit=False, auto_offset_reset='earliest', heartbeat_interval_ms=100) yield from consumer.start() yield from asyncio.sleep(0.2, loop=self.loop) yield from consumer.stop()
def test_manual_subscribe_pattern(self): msgs1 = yield from self.send_messages(0, range(0, 10)) msgs2 = yield from self.send_messages(1, range(10, 20)) available_msgs = msgs1 + msgs2 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*") yield from consumer.start() yield from consumer.seek_to_committed() result = [] for i in range(20): msg = yield from consumer.getone() result.append(msg.value) self.assertEqual(set(available_msgs), set(result)) yield from consumer.commit( {TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')}) yield from consumer.seek_to_committed(TopicPartition(self.topic, 0)) msg = yield from consumer.getone() self.assertEqual(msg.value, b'9') yield from consumer.commit( {TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')}) yield from 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,)) yield from consumer.start() yield from consumer.seek_to_committed() result = [] for i in range(10): msg = yield from consumer.getone() result.append(msg.value) self.assertEqual(set(msgs2), set(result)) self.assertEqual(consumer.subscription(), set([self.topic])) yield from consumer.stop()
def test_consumer_stop_cancels_pending_position_fetches(self): consumer = AIOKafkaConsumer(self.topic, loop=self.loop, bootstrap_servers=self.hosts, group_id='group-%s' % self.id()) yield from consumer.start() self.add_cleanup(consumer.stop) self.assertTrue(consumer._pending_position_fetches) pending_task = list(consumer._pending_position_fetches)[0] yield from consumer.stop() self.assertTrue(pending_task.cancelled())
def test_manual_subscribe_nogroup(self): msgs1 = yield from self.send_messages(0, range(0, 10)) msgs2 = yield from self.send_messages(1, range(10, 20)) available_msgs = msgs1 + msgs2 consumer = AIOKafkaConsumer( loop=self.loop, group_id=None, bootstrap_servers=self.hosts, auto_offset_reset='earliest', enable_auto_commit=False) consumer.subscribe(topics=(self.topic,)) yield from consumer.start() result = [] for i in range(20): msg = yield from consumer.getone() result.append(msg.value) self.assertEqual(set(available_msgs), set(result)) yield from consumer.stop()
def test_consumer_wait_topic(self): topic = "some-test-topic-for-autocreate" consumer = AIOKafkaConsumer( topic, loop=self.loop, bootstrap_servers=self.hosts) yield from consumer.start() consume_task = self.loop.create_task(consumer.getone()) # just to be sure getone does not fail (before produce) yield from asyncio.sleep(0.5, loop=self.loop) self.assertFalse(consume_task.done()) producer = AIOKafkaProducer( loop=self.loop, bootstrap_servers=self.hosts) yield from producer.start() yield from producer.send(topic, b'test msg') yield from producer.stop() data = yield from consume_task self.assertEqual(data.value, b'test msg') yield from consumer.stop()
def test_ssl_consume(self): # Produce by PLAINTEXT, Consume by SSL # Send 3 messages yield from self.send_messages(0, [1, 2, 3]) context = self.create_ssl_context() group = "group-{}".format(self.id()) consumer = AIOKafkaConsumer( self.topic, loop=self.loop, group_id=group, bootstrap_servers=[ "{}:{}".format(self.kafka_host, self.kafka_ssl_port)], enable_auto_commit=True, auto_offset_reset="earliest", security_protocol="SSL", ssl_context=context) yield from consumer.start() results = yield from consumer.getmany(timeout_ms=1000) [msgs] = results.values() # only 1 partition anyway msgs = [msg.value for msg in msgs] self.assertEqual(msgs, [b"1", b"2", b"3"]) yield from consumer.stop()
def test_producer_ssl(self): # Produce by SSL consume by PLAINTEXT topic = "test_ssl_produce" context = self.create_ssl_context() producer = AIOKafkaProducer( loop=self.loop, bootstrap_servers=[ "{}:{}".format(self.kafka_host, self.kafka_ssl_port)], security_protocol="SSL", ssl_context=context) yield from producer.start() yield from producer.send_and_wait(topic=topic, value=b"Super msg") yield from producer.stop() consumer = AIOKafkaConsumer( topic, loop=self.loop, bootstrap_servers=self.hosts, enable_auto_commit=True, auto_offset_reset="earliest") yield from consumer.start() msg = yield from consumer.getone() self.assertEqual(msg.value, b"Super msg") yield from consumer.stop()
def test_exclude_internal_topics(self): # Create random topic my_topic = "some_noninternal_topic" client = AIOKafkaClient( loop=self.loop, bootstrap_servers=self.hosts, client_id="test_autocreate") yield from client.bootstrap() yield from client._wait_on_metadata(my_topic) yield from client.close() # Check if only it will be subscribed pattern = "^.*$" consumer = AIOKafkaConsumer( loop=self.loop, bootstrap_servers=self.hosts, metadata_max_age_ms=200, group_id="some_group_1", auto_offset_reset="earliest", exclude_internal_topics=False) consumer.subscribe(pattern=pattern) yield from consumer.start() self.assertIn("__consumer_offsets", consumer.subscription()) yield from consumer._client.force_metadata_update() self.assertIn("__consumer_offsets", consumer.subscription()) yield from consumer.stop()
def test_consumer_rebalance_on_new_topic(self): # Test will create a consumer group and check if adding new topic # will trigger a group rebalance and assign partitions pattern = "^another-autocreate-pattern-.*$" client = AIOKafkaClient( loop=self.loop, bootstrap_servers=self.hosts, client_id="test_autocreate") yield from client.bootstrap() listener1 = StubRebalanceListener(loop=self.loop) listener2 = StubRebalanceListener(loop=self.loop) consumer1 = AIOKafkaConsumer( loop=self.loop, bootstrap_servers=self.hosts, metadata_max_age_ms=200, group_id="test-autocreate-rebalance", heartbeat_interval_ms=100) consumer1.subscribe(pattern=pattern, listener=listener1) yield from consumer1.start() consumer2 = AIOKafkaConsumer( loop=self.loop, bootstrap_servers=self.hosts, metadata_max_age_ms=200, group_id="test-autocreate-rebalance", heartbeat_interval_ms=100) consumer2.subscribe(pattern=pattern, listener=listener2) yield from consumer2.start() yield from asyncio.sleep(0.5, loop=self.loop) # bootstrap will take care of the initial group assignment self.assertEqual(consumer1.assignment(), set()) self.assertEqual(consumer2.assignment(), set()) listener1.reset() listener2.reset() # Lets force autocreation of a topic my_topic = "another-autocreate-pattern-1" yield from client._wait_on_metadata(my_topic) # Wait for group to stabilize assign1 = yield from listener1.wait_assign() assign2 = yield from listener2.wait_assign() # We expect 2 partitons for autocreated topics my_partitions = set([ TopicPartition(my_topic, 0), TopicPartition(my_topic, 1)]) self.assertEqual(assign1 | assign2, my_partitions) self.assertEqual( consumer1.assignment() | consumer2.assignment(), my_partitions) # Lets add another topic listener1.reset() listener2.reset() my_topic2 = "another-autocreate-pattern-2" yield from client._wait_on_metadata(my_topic2) # Wait for group to stabilize assign1 = yield from listener1.wait_assign() assign2 = yield from listener2.wait_assign() # We expect 2 partitons for autocreated topics my_partitions = set([ TopicPartition(my_topic, 0), TopicPartition(my_topic, 1), TopicPartition(my_topic2, 0), TopicPartition(my_topic2, 1)]) self.assertEqual(assign1 | assign2, my_partitions) self.assertEqual( consumer1.assignment() | consumer2.assignment(), my_partitions) yield from consumer1.stop() yield from consumer2.stop() yield from client.close()