Python AIOKafkaProducer.send_batch方法代码示例

# 需要导入模块: from aiokafka.producer import AIOKafkaProducer [as 别名]
# 或者: from aiokafka.producer.AIOKafkaProducer import send_batch [as 别名]
    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)
        yield from producer.start()

        partitions = yield from producer.partitions_for(self.topic)
        partition = partitions.pop()

        # silly method to find current offset for this partition
        resp = yield from 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
        while True:
            metadata = batch.append(key=key, value=value, timestamp=None)
            if metadata is None:
                break
            num += 1
        self.assertTrue(num > 0)
        self.assertEqual(batch.record_count(), num)

        # batch gets properly sent
        future = yield from producer.send_batch(
            batch, self.topic, partition=partition)
        resp = yield from future
        self.assertEqual(resp.topic, self.topic)
        self.assertEqual(resp.partition, partition)
        self.assertEqual(resp.offset, offset + 1)

        # batch accepts a too-large message if it's the first
        too_large = b'm' * (max_batch_size + 1)
        batch = producer.create_batch()
        metadata = batch.append(key=None, value=too_large, timestamp=None)
        self.assertIsNotNone(metadata)

        # batch rejects a too-large message if it's not the first
        batch = producer.create_batch()
        batch.append(key=None, value=b"short", timestamp=None)
        metadata = batch.append(key=None, value=too_large, timestamp=None)
        self.assertIsNone(metadata)
        yield from producer.stop()

        # batch can't be sent after closing time
        with self.assertRaises(ProducerClosed):
            yield from producer.send_batch(
                batch, self.topic, partition=partition)

# 需要导入模块: from aiokafka.producer import AIOKafkaProducer [as 别名]
# 或者: from aiokafka.producer.AIOKafkaProducer import send_batch [as 别名]
    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")