Python producer.AIOKafkaProducer类代码示例

    def test_producer_send_error(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts,
            retry_backoff_ms=100,
            linger_ms=5, request_timeout_ms=400)
        yield from producer.start()

        @asyncio.coroutine
        def mocked_send(nodeid, req):
            # RequestTimedOutCode error for partition=0
            return ProduceResponse[0]([(self.topic, [(0, 7, 0), (1, 0, 111)])])

        with mock.patch.object(producer.client, 'send') as mocked:
            mocked.side_effect = mocked_send
            fut1 = yield from producer.send(self.topic, b'text1', partition=0)
            fut2 = yield from producer.send(self.topic, b'text2', partition=1)
            with self.assertRaises(RequestTimedOutError):
                yield from fut1
            resp = yield from fut2
            self.assertEqual(resp.offset, 111)

        @asyncio.coroutine
        def mocked_send_with_sleep(nodeid, req):
            # RequestTimedOutCode error for partition=0
            yield from asyncio.sleep(0.1, loop=self.loop)
            return ProduceResponse[0]([(self.topic, [(0, 7, 0)])])

        with mock.patch.object(producer.client, 'send') as mocked:
            mocked.side_effect = mocked_send_with_sleep
            with self.assertRaises(RequestTimedOutError):
                future = yield from producer.send(
                    self.topic, b'text1', partition=0)
                yield from future
        yield from producer.stop()

    def send_messages(self, partition, messages, *, topic=None,
                      timestamp_ms=None, return_inst=False, headers=None):
        topic = topic or self.topic
        ret = []
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        yield from producer.start()
        try:
            yield from self.wait_topic(producer.client, topic)

            for msg in messages:
                if isinstance(msg, str):
                    msg = msg.encode()
                elif isinstance(msg, int):
                    msg = str(msg).encode()
                future = yield from producer.send(
                    topic, msg, partition=partition,
                    timestamp_ms=timestamp_ms, headers=headers)
                resp = yield from future
                self.assertEqual(resp.topic, topic)
                self.assertEqual(resp.partition, partition)
                if return_inst:
                    ret.append(resp)
                else:
                    ret.append(msg)
        finally:
            yield from producer.stop()
        return ret

    def test_producer_correct_time_returned(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        send_time = (time.time() * 1000)
        res = yield from producer.send_and_wait(
            "XXXX", b'text1', partition=0)
        self.assertLess(res.timestamp - send_time, 1000)  # 1s

        res = yield from producer.send_and_wait(
            "XXXX", b'text1', partition=0, timestamp_ms=123123123)
        self.assertEqual(res.timestamp, 123123123)

        expected_timestamp = 999999999

        @asyncio.coroutine
        def mocked_send(*args, **kw):
            # There's no easy way to set LOG_APPEND_TIME on server, so use this
            # hack for now.
            return ProduceResponse[2](
                topics=[
                    ('XXXX', [(0, 0, 0, expected_timestamp)])],
                throttle_time_ms=0)

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

            res = yield from producer.send_and_wait(
                "XXXX", b'text1', partition=0)
            self.assertEqual(res.timestamp_type, LOG_APPEND_TIME)
            self.assertEqual(res.timestamp, expected_timestamp)

 def test_producer_send_noack(self):
     producer = AIOKafkaProducer(loop=self.loop, bootstrap_servers=self.hosts, acks=0)
     yield from producer.start()
     yield from self.wait_topic(producer.client, self.topic)
     fut1 = yield from producer.send(self.topic, b"hello, Kafka!", partition=0)
     fut2 = yield from producer.send(self.topic, b"hello, Kafka!", partition=1)
     done, _ = yield from asyncio.wait([fut1, fut2], loop=self.loop)
     for item in done:
         self.assertEqual(item.result(), None)

    async def test_producer_transactional_flush_2_batches_before_commit(self):
        # We need to be sure if batches that are pending and batches that are
        # queued will be waited. To test this we need at least 2 batches on
        # the same partition. They will be sent one at a time.

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

        await producer.begin_transaction()

        batch = producer.create_batch()
        batch.append(timestamp=None, key=None, value=b"1")
        batch.close()
        fut1 = await producer.send_batch(batch, self.topic, partition=0)

        batch = producer.create_batch()
        batch.append(timestamp=None, key=None, value=b"2")
        batch.close()
        fut2 = await producer.send_batch(batch, self.topic, partition=0)

        await producer.commit_transaction()

        self.assertTrue(fut1.done())
        self.assertTrue(fut2.done())

    def test_producer_api_version(self):
        for text_version, api_version in [
                ("auto", (0, 9, 0)),
                ("0.9.1", (0, 9, 1)),
                ("0.10.0", (0, 10, 0)),
                ("0.11", (0, 11, 0)),
                ("0.12.1", (0, 12, 1)),
                ("1.0.2", (1, 0, 2))]:
            producer = AIOKafkaProducer(
                loop=self.loop, bootstrap_servers=self.hosts,
                api_version=text_version)
            self.assertEqual(producer.client.api_version, api_version)
            yield from producer.stop()

        # invalid cases
        for version in ["0", "1", "0.10.0.1"]:
            with self.assertRaises(ValueError):
                AIOKafkaProducer(
                    loop=self.loop, bootstrap_servers=self.hosts,
                    api_version=version)
        for version in [(0, 9), (0, 9, 1)]:
            with self.assertRaises(TypeError):
                AIOKafkaProducer(
                    loop=self.loop, bootstrap_servers=self.hosts,
                    api_version=version)

    async def test_producer_transactional_simple(self):
        # The test here will just check if we can do simple produce with
        # transactional_id option and minimal setup.

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

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

        consumer = AIOKafkaConsumer(
            self.topic, loop=self.loop,
            bootstrap_servers=self.hosts,
            auto_offset_reset="earliest")
        await consumer.start()
        self.add_cleanup(consumer.stop)
        msg = await 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)

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

        batch = producer.create_batch()
        batch.append(timestamp=None, key=None, value=b"2")
        batch.close()

        # Can not send if not yet in transaction
        with self.assertRaises(IllegalOperation):
            await producer.send_batch(batch, self.topic, partition=0)

        await producer.begin_transaction()
        await producer.send(self.topic, value=b"1", partition=0)
        commit_task = ensure_future(producer.commit_transaction())
        await asyncio.sleep(0.001, loop=self.loop)
        self.assertFalse(commit_task.done())

        # Already not in transaction
        with self.assertRaises(IllegalOperation):
            await producer.send_batch(batch, self.topic, partition=0)

        await commit_task
        # Transaction needs to be restarted
        with self.assertRaises(IllegalOperation):
            await producer.send_batch(batch, self.topic, partition=0)

 def test_producer_indempotence_not_supported(self):
     producer = AIOKafkaProducer(
         loop=self.loop, bootstrap_servers=self.hosts,
         enable_idempotence=True)
     producer
     with self.assertRaises(UnsupportedVersionError):
         yield from producer.start()
     yield from producer.stop()

    def test_producer_send_with_headers_raise_error(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        with self.assertRaises(UnsupportedVersionError):
            yield from producer.send(
                self.topic, b'msg', partition=0,
                headers=[("type", b"Normal")])

    def test_producer_send_with_headers(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        yield from producer.start()
        self.add_cleanup(producer.stop)

        fut = yield from producer.send(
            self.topic, b'msg', partition=0, headers=[("type", b"Normal")])
        resp = yield from fut
        self.assertEqual(resp.partition, 0)

    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)

    def test_producer_send(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        yield from producer.start()
        yield from self.wait_topic(producer.client, self.topic)
        with self.assertRaisesRegexp(AssertionError, 'value must be bytes'):
            yield from producer.send(self.topic, 'hello, Kafka!')
        future = yield from producer.send(self.topic, b'hello, Kafka!')
        resp = yield from future
        self.assertEqual(resp.topic, self.topic)
        self.assertTrue(resp.partition in (0, 1))
        self.assertEqual(resp.offset, 0)

        fut = yield from producer.send(self.topic, b'second msg', partition=1)
        resp = yield from fut
        self.assertEqual(resp.partition, 1)

        future = yield from producer.send(self.topic, b'value', key=b'KEY')
        resp = yield from future
        self.assertTrue(resp.partition in (0, 1))

        resp = yield from producer.send_and_wait(self.topic, b'value')
        self.assertTrue(resp.partition in (0, 1))

        yield from producer.stop()
        with self.assertRaises(ProducerClosed):
            yield from producer.send(self.topic, b'value', key=b'KEY')

    def test_producer_send(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        yield from producer.start()
        self.add_cleanup(producer.stop)
        with self.assertRaises(TypeError):
            yield from producer.send(self.topic, 'hello, Kafka!', partition=0)
        future = yield from producer.send(
            self.topic, b'hello, Kafka!', partition=0)
        resp = yield from future
        self.assertEqual(resp.topic, self.topic)
        self.assertTrue(resp.partition in (0, 1))
        self.assertEqual(resp.offset, 0)

        fut = yield from producer.send(self.topic, b'second msg', partition=1)
        resp = yield from fut
        self.assertEqual(resp.partition, 1)

        future = yield from producer.send(self.topic, b'value', key=b'KEY')
        resp = yield from future
        self.assertTrue(resp.partition in (0, 1))

        resp = yield from producer.send_and_wait(self.topic, b'value')
        self.assertTrue(resp.partition in (0, 1))

        yield from producer.stop()
        with self.assertRaises(ProducerClosed):
            yield from producer.send(self.topic, b'value', key=b'KEY')

    def test_producer_warn_unclosed(self):
        producer = AIOKafkaProducer(
            loop=self.loop, bootstrap_servers=self.hosts)
        producer_ref = weakref.ref(producer)
        yield from producer.start()

        with self.silence_loop_exception_handler():
            with self.assertWarnsRegex(
                    ResourceWarning, "Unclosed AIOKafkaProducer"):
                del producer
                gc.collect()
        # Assure that the reference was properly collected
        self.assertIsNone(producer_ref())