Exemple #1
0
    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()
Exemple #2
0
    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()
Exemple #3
0
 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()
Exemple #4
0
    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()
Exemple #5
0
    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)
        consumer.subscribe(topics=('some_topic_unknown',))
        with self.assertRaises(UnknownTopicOrPartitionError):
            yield from consumer.start()

        with self.assertRaises(UnknownTopicOrPartitionError):
            yield from consumer.assign([TopicPartition(self.topic, 2222)])
Exemple #6
0
    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())
Exemple #7
0
 def consumer_factory(self, **kwargs):
     enable_auto_commit = kwargs.pop('enable_auto_commit', True)
     auto_offset_reset = kwargs.pop('auto_offset_reset', 'earliest')
     group = kwargs.pop('group', 'group-%s' % self.id())
     consumer = AIOKafkaConsumer(
         self.topic, loop=self.loop, group_id=group,
         bootstrap_servers=self.hosts,
         enable_auto_commit=enable_auto_commit,
         auto_offset_reset=auto_offset_reset,
         **kwargs)
     yield from consumer.start()
     if group is not None:
         yield from consumer.seek_to_committed()
     return consumer
Exemple #8
0
 def consumer_factory(self, **kwargs):
     enable_auto_commit = kwargs.pop('enable_auto_commit', True)
     auto_offset_reset = kwargs.pop('auto_offset_reset', 'earliest')
     group = kwargs.pop('group', 'group-%s' % self.id())
     consumer = AIOKafkaConsumer(
         self.topic, loop=self.loop, group_id=group,
         bootstrap_servers=self.hosts,
         enable_auto_commit=enable_auto_commit,
         auto_offset_reset=auto_offset_reset,
         **kwargs)
     yield from consumer.start()
     if group is not None:
         yield from consumer.seek_to_committed()
     return consumer
Exemple #9
0
    def test_offset_reset_manual(self):
        yield from self.send_messages(0, list(range(0, 10)))

        consumer = AIOKafkaConsumer(
            self.topic,
            loop=self.loop, bootstrap_servers=self.hosts,
            metadata_max_age_ms=200, group_id="offset_reset_group",
            auto_offset_reset="none")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)

        with self.assertRaises(OffsetOutOfRangeError):
            yield from consumer.getmany(timeout_ms=1000)

        with self.assertRaises(OffsetOutOfRangeError):
            yield from consumer.getone()
Exemple #10
0
    def test_consumer_seek_on_unassigned(self):
        tp0 = TopicPartition(self.topic, 0)
        tp1 = TopicPartition(self.topic, 1)
        consumer = AIOKafkaConsumer(loop=self.loop,
                                    group_id=None,
                                    bootstrap_servers=self.hosts)
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        consumer.assign([tp0])

        with self.assertRaises(IllegalStateError):
            yield from consumer.seek_to_beginning(tp1)
        with self.assertRaises(IllegalStateError):
            yield from consumer.seek_to_committed(tp1)
        with self.assertRaises(IllegalStateError):
            yield from consumer.seek_to_end(tp1)
Exemple #11
0
    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()
Exemple #12
0
    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()
Exemple #13
0
    def test_consumer_subscribe_pattern_autocreate_no_group_id(self):
        pattern = "^no-group-pattern-.*$"
        consumer = AIOKafkaConsumer(loop=self.loop,
                                    bootstrap_servers=self.hosts,
                                    metadata_max_age_ms=200,
                                    group_id=None,
                                    fetch_max_wait_ms=50,
                                    auto_offset_reset="earliest")
        self.add_cleanup(consumer.stop)
        yield from consumer.start()
        consumer.subscribe(pattern=pattern)
        # Start getter for the topics. Should not create any topics
        consume_task = self.loop.create_task(consumer.getone())
        yield from asyncio.sleep(0.3, loop=self.loop)
        self.assertFalse(consume_task.done())
        self.assertEqual(consumer.subscription(), set())

        # Now lets autocreate the topic by fetching metadata for it.
        producer = AIOKafkaProducer(loop=self.loop,
                                    bootstrap_servers=self.hosts)
        self.add_cleanup(producer.stop)
        yield from producer.start()
        my_topic = "no-group-pattern-1"
        yield from producer.client._wait_on_metadata(my_topic)
        # Wait for consumer to refresh metadata with new topic
        yield from asyncio.sleep(0.3, loop=self.loop)
        self.assertFalse(consume_task.done())
        self.assertTrue(consumer._client.cluster.topics() >= {my_topic})
        self.assertEqual(consumer.subscription(), {my_topic})

        # Add another topic
        my_topic2 = "no-group-pattern-2"
        yield from producer.client._wait_on_metadata(my_topic2)
        # Wait for consumer to refresh metadata with new topic
        yield from asyncio.sleep(0.3, loop=self.loop)
        self.assertFalse(consume_task.done())
        self.assertTrue(
            consumer._client.cluster.topics() >= {my_topic, my_topic2})
        self.assertEqual(consumer.subscription(), {my_topic, my_topic2})

        # Now lets actualy produce some data and verify that it is consumed
        yield from producer.send(my_topic, b'test msg')
        data = yield from consume_task
        self.assertEqual(data.value, b'test msg')
Exemple #14
0
    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()
Exemple #15
0
    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()
Exemple #16
0
    def test_producer_indempotence_no_duplicates(self):
        # Indempotent producer should retry produce in case of timeout error
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            enable_idempotence=True,
            request_timeout_ms=2000)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        original_send = producer.client.send
        retry = [0]

        @asyncio.coroutine
        def mocked_send(*args, **kw):
            result = yield from original_send(*args, **kw)
            if result.API_KEY == ProduceResponse[0].API_KEY and retry[0] < 2:
                retry[0] += 1
                raise RequestTimedOutError
            return result

        with mock.patch.object(producer.client, 'send') as mocked:
            mocked.side_effect = mocked_send

            meta = yield from producer.send_and_wait(
                self.topic, b'hello, Kafka!')

        consumer = AIOKafkaConsumer(
            self.topic, loop=self.loop,
            bootstrap_servers=self.hosts,
            auto_offset_reset="earliest")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        msg = yield from consumer.getone()
        self.assertEqual(msg.offset, meta.offset)
        self.assertEqual(msg.timestamp, meta.timestamp)
        self.assertEqual(msg.value, b"hello, Kafka!")
        self.assertEqual(msg.key, None)

        with self.assertRaises(asyncio.TimeoutError):
            yield from asyncio.wait_for(consumer.getone(), timeout=0.5)
Exemple #17
0
    def test_producer_indempotence_no_duplicates(self):
        # Indempotent producer should retry produce in case of timeout error
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            enable_idempotence=True,
            request_timeout_ms=2000)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        original_send = producer.client.send
        retry = [0]

        @asyncio.coroutine
        def mocked_send(*args, **kw):
            result = yield from original_send(*args, **kw)
            if result.API_KEY == ProduceResponse[0].API_KEY and retry[0] < 2:
                retry[0] += 1
                raise RequestTimedOutError
            return result

        with mock.patch.object(producer.client, 'send') as mocked:
            mocked.side_effect = mocked_send

            meta = yield from producer.send_and_wait(
                self.topic, b'hello, Kafka!')

        consumer = AIOKafkaConsumer(
            self.topic, loop=self.loop,
            bootstrap_servers=self.hosts,
            auto_offset_reset="earliest")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        msg = yield from consumer.getone()
        self.assertEqual(msg.offset, meta.offset)
        self.assertEqual(msg.timestamp, meta.timestamp)
        self.assertEqual(msg.value, b"hello, Kafka!")
        self.assertEqual(msg.key, None)

        with self.assertRaises(asyncio.TimeoutError):
            yield from asyncio.wait_for(consumer.getone(), timeout=0.5)
Exemple #18
0
    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()
Exemple #19
0
    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()
Exemple #20
0
    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()
Exemple #21
0
    def test_producer_indempotence_simple(self):
        # The test here will just check if we can do simple produce with
        # enable_idempotence option, as no specific API changes is expected.

        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            enable_idempotence=True)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        meta = yield from producer.send_and_wait(self.topic, b'hello, Kafka!')

        consumer = AIOKafkaConsumer(
            self.topic, loop=self.loop,
            bootstrap_servers=self.hosts,
            auto_offset_reset="earliest")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        msg = yield from consumer.getone()
        self.assertEqual(msg.offset, meta.offset)
        self.assertEqual(msg.timestamp, meta.timestamp)
        self.assertEqual(msg.value, b"hello, Kafka!")
        self.assertEqual(msg.key, None)
Exemple #22
0
    def test_producer_indempotence_simple(self):
        # The test here will just check if we can do simple produce with
        # enable_idempotence option, as no specific API changes is expected.

        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            enable_idempotence=True)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        meta = yield from producer.send_and_wait(self.topic, b'hello, Kafka!')

        consumer = AIOKafkaConsumer(
            self.topic, loop=self.loop,
            bootstrap_servers=self.hosts,
            auto_offset_reset="earliest")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        msg = yield from consumer.getone()
        self.assertEqual(msg.offset, meta.offset)
        self.assertEqual(msg.timestamp, meta.timestamp)
        self.assertEqual(msg.value, b"hello, Kafka!")
        self.assertEqual(msg.key, None)
Exemple #23
0
    def test_producer_leader_change_preserves_order(self):
        # Before 0.5.0 we did not lock partition until a response came from
        # the server, but locked the node itself.
        # For example: Say the sender sent a request to node 1 and before an
        # failure answer came we updated metadata and leader become node 0.
        # This way we may send the next batch to node 0 without waiting for
        # node 1 batch to be reenqueued, resulting in out-of-order batches

        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts, linger_ms=1000)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        # Alter metadata to convince the producer, that leader or partition 0
        # is a different node
        yield from producer.partitions_for(self.topic)
        topic_meta = producer._metadata._partitions[self.topic]
        real_leader = topic_meta[0].leader
        topic_meta[0] = topic_meta[0]._replace(leader=real_leader + 1)

        # Make sure the first request for produce takes more time
        original_send = producer.client.send

        @asyncio.coroutine
        def mocked_send(node_id, request, *args, **kw):
            if node_id != real_leader and \
                    request.API_KEY == ProduceResponse[0].API_KEY:
                yield from asyncio.sleep(2, loop=self.loop)

            result = yield from original_send(node_id, request, *args, **kw)
            return result
        producer.client.send = mocked_send

        # Send Batch 1. This will end up waiting for some time on fake leader
        batch = producer.create_batch()
        meta = batch.append(key=b"key", value=b"1", timestamp=None)
        batch.close()
        fut = yield from producer.send_batch(
            batch, self.topic, partition=0)

        # Make sure we sent the request
        yield from asyncio.sleep(0.1, loop=self.loop)
        # Update metadata to return leader to real one
        yield from producer.client.force_metadata_update()

        # Send Batch 2, that if it's bugged will go straight to the real node
        batch2 = producer.create_batch()
        meta2 = batch2.append(key=b"key", value=b"2", timestamp=None)
        batch2.close()
        fut2 = yield from producer.send_batch(
            batch2, self.topic, partition=0)

        batch_meta = yield from fut
        batch_meta2 = yield from fut2

        # Check the order of messages
        consumer = AIOKafkaConsumer(
            self.topic, loop=self.loop,
            bootstrap_servers=self.hosts,
            auto_offset_reset="earliest")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        msg = yield from consumer.getone()
        self.assertEqual(msg.offset, batch_meta.offset)
        self.assertEqual(msg.timestamp or -1, meta.timestamp)
        self.assertEqual(msg.value, b"1")
        self.assertEqual(msg.key, b"key")
        msg2 = yield from consumer.getone()
        self.assertEqual(msg2.offset, batch_meta2.offset)
        self.assertEqual(msg2.timestamp or -1, meta2.timestamp)
        self.assertEqual(msg2.value, b"2")
        self.assertEqual(msg2.key, b"key")
Exemple #24
0
    def test_rebalance_listener_with_coroutines(self):
        yield from self.send_messages(0, list(range(0, 10)))
        yield from self.send_messages(1, list(range(10, 20)))

        main_self = self

        class SimpleRebalanceListener(ConsumerRebalanceListener):
            def __init__(self, consumer):
                self.consumer = consumer
                self.revoke_mock = mock.Mock()
                self.assign_mock = mock.Mock()

            @asyncio.coroutine
            def on_partitions_revoked(self, revoked):
                self.revoke_mock(revoked)
                # If this commit would fail we will end up with wrong msgs
                # eturned in test below
                yield from self.consumer.commit()
                # Confirm that coordinator is actually waiting for callback to
                # complete
                yield from asyncio.sleep(0.2, loop=main_self.loop)
                main_self.assertTrue(
                    self.consumer._coordinator.needs_join_prepare)

            @asyncio.coroutine
            def on_partitions_assigned(self, assigned):
                self.assign_mock(assigned)
                # Confirm that coordinator is actually waiting for callback to
                # complete
                yield from asyncio.sleep(0.2, loop=main_self.loop)
                main_self.assertFalse(
                    self.consumer._coordinator.needs_join_prepare)

        tp0 = TopicPartition(self.topic, 0)
        tp1 = TopicPartition(self.topic, 1)
        consumer1 = AIOKafkaConsumer(
            loop=self.loop,
            group_id="test_rebalance_listener_with_coroutines",
            bootstrap_servers=self.hosts,
            enable_auto_commit=False,
            auto_offset_reset="earliest")
        listener1 = SimpleRebalanceListener(consumer1)
        consumer1.subscribe([self.topic], listener=listener1)
        yield from consumer1.start()
        self.add_cleanup(consumer1.stop)

        msg = yield from consumer1.getone(tp0)
        self.assertEqual(msg.value, b"0")
        msg = yield from consumer1.getone(tp1)
        self.assertEqual(msg.value, b"10")
        listener1.revoke_mock.assert_called_with(set([]))
        listener1.assign_mock.assert_called_with(set([tp0, tp1]))

        # By adding a 2nd consumer we trigger rebalance
        consumer2 = AIOKafkaConsumer(
            loop=self.loop,
            group_id="test_rebalance_listener_with_coroutines",
            bootstrap_servers=self.hosts,
            enable_auto_commit=False,
            auto_offset_reset="earliest")
        listener2 = SimpleRebalanceListener(consumer2)
        consumer2.subscribe([self.topic], listener=listener2)
        yield from consumer2.start()
        self.add_cleanup(consumer2.stop)

        msg1 = yield from consumer1.getone()
        msg2 = yield from consumer2.getone()
        # We can't predict the assignment in test
        if consumer1.assignment() == set([tp1]):
            msg1, msg2 = msg2, msg1
            c1_assignment = set([tp1])
            c2_assignment = set([tp0])
        else:
            c1_assignment = set([tp0])
            c2_assignment = set([tp1])

        self.assertEqual(msg1.value, b"1")
        self.assertEqual(msg2.value, b"11")

        listener1.revoke_mock.assert_called_with(set([tp0, tp1]))
        self.assertEqual(listener1.revoke_mock.call_count, 2)
        listener1.assign_mock.assert_called_with(c1_assignment)
        self.assertEqual(listener1.assign_mock.call_count, 2)

        listener2.revoke_mock.assert_called_with(set([]))
        self.assertEqual(listener2.revoke_mock.call_count, 1)
        listener2.assign_mock.assert_called_with(c2_assignment)
        self.assertEqual(listener2.assign_mock.call_count, 1)
Exemple #25
0
    def test_producer_leader_change_preserves_order(self):
        # Before 0.5.0 we did not lock partition until a response came from
        # the server, but locked the node itself.
        # For example: Say the sender sent a request to node 1 and before an
        # failure answer came we updated metadata and leader become node 0.
        # This way we may send the next batch to node 0 without waiting for
        # node 1 batch to be reenqueued, resulting in out-of-order batches

        producer = AIOKafkaProducer(loop=self.loop,
                                    bootstrap_servers=self.hosts,
                                    linger_ms=1000)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        # Alter metadata to convince the producer, that leader or partition 0
        # is a different node
        yield from producer.partitions_for(self.topic)
        topic_meta = producer._metadata._partitions[self.topic]
        real_leader = topic_meta[0].leader
        topic_meta[0] = topic_meta[0]._replace(leader=real_leader + 1)

        # Make sure the first request for produce takes more time
        original_send = producer.client.send

        @asyncio.coroutine
        def mocked_send(node_id, request, *args, **kw):
            if node_id != real_leader and \
                    request.API_KEY == ProduceResponse[0].API_KEY:
                yield from asyncio.sleep(2, loop=self.loop)

            result = yield from original_send(node_id, request, *args, **kw)
            return result

        producer.client.send = mocked_send

        # Send Batch 1. This will end up waiting for some time on fake leader
        batch = producer.create_batch()
        meta = batch.append(key=b"key", value=b"1", timestamp=None)
        batch.close()
        fut = yield from producer.send_batch(batch, self.topic, partition=0)

        # Make sure we sent the request
        yield from asyncio.sleep(0.1, loop=self.loop)
        # Update metadata to return leader to real one
        yield from producer.client.force_metadata_update()

        # Send Batch 2, that if it's bugged will go straight to the real node
        batch2 = producer.create_batch()
        meta2 = batch2.append(key=b"key", value=b"2", timestamp=None)
        batch2.close()
        fut2 = yield from producer.send_batch(batch2, self.topic, partition=0)

        batch_meta = yield from fut
        batch_meta2 = yield from fut2

        # Check the order of messages
        consumer = AIOKafkaConsumer(self.topic,
                                    loop=self.loop,
                                    bootstrap_servers=self.hosts,
                                    auto_offset_reset="earliest")
        yield from consumer.start()
        self.add_cleanup(consumer.stop)
        msg = yield from consumer.getone()
        self.assertEqual(msg.offset, batch_meta.offset)
        self.assertEqual(msg.timestamp or -1, meta.timestamp)
        self.assertEqual(msg.value, b"1")
        self.assertEqual(msg.key, b"key")
        msg2 = yield from consumer.getone()
        self.assertEqual(msg2.offset, batch_meta2.offset)
        self.assertEqual(msg2.timestamp or -1, meta2.timestamp)
        self.assertEqual(msg2.value, b"2")
        self.assertEqual(msg2.key, b"key")
Exemple #26
0
    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()