How to use aiokafka - 10 common examples

async def test_coordinator__maybe_do_autocommit(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        subscription = SubscriptionState(loop=self.loop)
        tp = TopicPartition("topic1", 0)
        coordinator = GroupCoordinator(
            client, subscription, loop=self.loop,
            heartbeat_interval_ms=20000, auto_commit_interval_ms=1000,
            retry_backoff_ms=50)
        coordinator._coordination_task.cancel()  # disable for test
        try:
            await coordinator._coordination_task
        except asyncio.CancelledError:
            pass
        coordinator._coordination_task = self.loop.create_task(
            asyncio.sleep(0.1, loop=self.loop)
        )
        self.add_cleanup(coordinator.close)

        coordinator._do_commit_offsets = mocked = mock.Mock()
        loop = self.loop

async def test_producer_send_batch(self):
        key = b'test key'
        value = b'test value'
        max_batch_size = 10000

        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            max_batch_size=max_batch_size)
        await producer.start()

        partitions = await producer.partitions_for(self.topic)
        partition = partitions.pop()

        # silly method to find current offset for this partition
        resp = await producer.send_and_wait(
            self.topic, value=b'discovering offset', partition=partition)
        offset = resp.offset

        # only fills up to its limits, then returns None
        batch = producer.create_batch()
        self.assertEqual(batch.record_count(), 0)
        num = 0

async def test_producer_transactional_fences_off_previous(self):
        # Test 2 producers fencing one another by using the same
        # transactional_id

        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            transactional_id="sobaka_producer", client_id="p1")
        await producer.start()
        self.add_cleanup(producer.stop)

        producer2 = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            transactional_id="sobaka_producer", client_id="p2")
        await producer2.start()
        self.add_cleanup(producer2.stop)
        async with producer2.transaction():
            await producer2.send_and_wait(self.topic, b'hello, Kafka! 2')

        with self.assertRaises(ProducerFenced):
            async with producer.transaction():
                await producer.send_and_wait(self.topic, b'hello, Kafka!')

async def test_manual_subscribe_pattern(self):
        msgs1 = await self.send_messages(0, range(0, 10))
        msgs2 = await 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*")
        await consumer.start()
        self.add_cleanup(consumer.stop)
        await consumer.seek_to_committed()
        result = []
        for i in range(20):
            msg = await consumer.getone()
            result.append(msg.value)
        self.assertEqual(set(available_msgs), set(result))

        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
        await consumer.seek_to_committed(TopicPartition(self.topic, 0))

for i in range(20):
            msg = await consumer.getone()
            result.append(msg.value)
        self.assertEqual(set(available_msgs), set(result))

        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
        await consumer.seek_to_committed(TopicPartition(self.topic, 0))
        msg = await consumer.getone(TopicPartition(self.topic, 0))
        self.assertEqual(msg.value, b'9')
        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
        await 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,))
        await consumer.start()
        self.add_cleanup(consumer.stop)
        await consumer.seek_to_committed()
        result = []
        for i in range(10):
            msg = await consumer.getone()
            result.append(msg.value)
        self.assertEqual(set(msgs2), set(result))
        self.assertEqual(consumer.subscription(), set([self.topic]))

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

        with mock.patch.object(producer._sender, '_send_produce_req') as m:
            m.side_effect = KeyError

            with self.assertRaisesRegex(
                    KafkaError, "Unexpected error during batch delivery"):
                await producer.send_and_wait(
                    self.topic, b'hello, Kafka!')

        with self.assertRaisesRegex(
                KafkaError, "Unexpected error during batch delivery"):
            await producer.send_and_wait(
                self.topic, b'hello, Kafka!')

async def test_coordinator__maybe_do_autocommit(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        subscription = SubscriptionState(loop=self.loop)
        tp = TopicPartition("topic1", 0)
        coordinator = GroupCoordinator(
            client, subscription, loop=self.loop,
            heartbeat_interval_ms=20000, auto_commit_interval_ms=1000,
            retry_backoff_ms=50)
        coordinator._coordination_task.cancel()  # disable for test
        try:
            await coordinator._coordination_task
        except asyncio.CancelledError:
            pass
        coordinator._coordination_task = self.loop.create_task(
            asyncio.sleep(0.1, loop=self.loop)
        )
        self.add_cleanup(coordinator.close)

        coordinator._do_commit_offsets = mocked = mock.Mock()
        loop = self.loop

await consumer.start()
        self.add_cleanup(consumer.stop)
        await consumer.seek_to_committed()
        result = []
        for i in range(20):
            msg = await consumer.getone()
            result.append(msg.value)
        self.assertEqual(set(available_msgs), set(result))

        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
        await consumer.seek_to_committed(TopicPartition(self.topic, 0))
        msg = await consumer.getone(TopicPartition(self.topic, 0))
        self.assertEqual(msg.value, b'9')
        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
        await 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,))
        await consumer.start()
        self.add_cleanup(consumer.stop)
        await consumer.seek_to_committed()
        result = []
        for i in range(10):
            msg = await consumer.getone()
            result.append(msg.value)
        self.assertEqual(set(msgs2), set(result))

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*")
        await consumer.start()
        self.add_cleanup(consumer.stop)
        await consumer.seek_to_committed()
        result = []
        for i in range(20):
            msg = await consumer.getone()
            result.append(msg.value)
        self.assertEqual(set(available_msgs), set(result))

        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
        await consumer.seek_to_committed(TopicPartition(self.topic, 0))
        msg = await consumer.getone(TopicPartition(self.topic, 0))
        self.assertEqual(msg.value, b'9')
        await consumer.commit(
            {TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
        await 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,))
        await consumer.start()
        self.add_cleanup(consumer.stop)
        await consumer.seek_to_committed()

async def test_consumer_commit_validation(self):
        consumer = await self.consumer_factory()
        self.add_cleanup(consumer.stop)

        tp = TopicPartition(self.topic, 0)
        offset = await consumer.position(tp)
        offset_and_metadata = OffsetAndMetadata(offset, "")

        with self.assertRaises(ValueError):
            await consumer.commit({})
        with self.assertRaises(ValueError):
            await consumer.commit("something")
        with self.assertRaises(ValueError):
            await consumer.commit({tp: (offset, "metadata", 100)})
        with self.assertRaisesRegex(
                ValueError, "Key should be TopicPartition instance"):
            await consumer.commit({"my_topic": offset_and_metadata})
        with self.assertRaisesRegex(
                ValueError, "Metadata should be a string"):
            await consumer.commit({tp: (offset, 1000)})
        with self.assertRaisesRegex(
                ValueError, "Metadata should be a string"):
            await consumer.commit({tp: (offset, b"\x00\x02")})