asyncio_test.py

#!/usr/bin/env python3
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

"""
Unit tests for asyncio Pulsar client API.
"""

# pylint: disable=missing-function-docstring

import asyncio
import time
from typing import List
from unittest import (
    main,
    IsolatedAsyncioTestCase,
)

import pulsar  # pylint: disable=import-error
from pulsar.asyncio import (  # pylint: disable=import-error
    Client,
    Consumer,
    Producer,
    PulsarException,
)
from pulsar.schema import (  # pylint: disable=import-error
    AvroSchema,
    Integer,
    Record,
    String,
)

SERVICE_URL = 'pulsar://localhost:6650'

class AsyncioTest(IsolatedAsyncioTestCase):
    """Test cases for asyncio Pulsar client."""

    async def asyncSetUp(self) -> None:
        self._client = Client(SERVICE_URL,
                              operation_timeout_seconds=5)

    async def asyncTearDown(self) -> None:
        await self._client.close()

    async def test_batch_end_to_end(self):
        topic = f'asyncio-test-batch-e2e-{time.time()}'
        producer = await self._client.create_producer(topic,
                                                      producer_name="my-producer")
        self.assertEqual(producer.topic(), f'persistent://public/default/{topic}')
        self.assertEqual(producer.producer_name(), "my-producer")
        tasks = []
        for i in range(5):
            tasks.append(asyncio.create_task(producer.send(f'msg-{i}'.encode())))
        msg_ids = await asyncio.gather(*tasks)
        self.assertEqual(len(msg_ids), 5)
        # pylint: disable=fixme
        # TODO: the result is wrong due to https://github.com/apache/pulsar-client-cpp/issues/531
        self.assertEqual(producer.last_sequence_id(), 8)
        ledger_id = msg_ids[0].ledger_id()
        entry_id = msg_ids[0].entry_id()
        # These messages should be in the same entry
        for i in range(5):
            msg_id = msg_ids[i]
            print(f'{i} was sent to {msg_id}')
            self.assertIsInstance(msg_id, pulsar.MessageId)
            self.assertEqual(msg_ids[i].ledger_id(), ledger_id)
            self.assertEqual(msg_ids[i].entry_id(), entry_id)
            self.assertEqual(msg_ids[i].batch_index(), i)

        consumer = await self._client.subscribe(topic, 'sub',
                                                initial_position=pulsar.InitialPosition.Earliest)
        for i in range(5):
            msg = await consumer.receive()
            self.assertEqual(msg.data(), f'msg-{i}'.encode())
        await consumer.close()

        # create a different subscription to verify initial position is latest by default
        consumer = await self._client.subscribe(topic, 'sub2')
        await producer.send(b'final-message')
        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'final-message')

    async def test_send_keyed_message(self):
        topic = f'asyncio-test-send-keyed-message-{time.time()}'
        producer = await self._client.create_producer(topic)
        consumer = await self._client.subscribe(topic, 'sub')
        await producer.send(b'msg', partition_key='key0',
                                     ordering_key="key1", properties={'my-prop': 'my-value'})

        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'msg')
        self.assertEqual(msg.partition_key(), 'key0')
        self.assertEqual(msg.ordering_key(), 'key1')
        self.assertEqual(msg.properties(), {'my-prop': 'my-value'})

    async def test_flush(self):
        topic = f'asyncio-test-flush-{time.time()}'
        producer = await self._client.create_producer(topic, batching_max_messages=3,
                                                      batching_max_publish_delay_ms=60000)
        tasks = []
        tasks.append(asyncio.create_task(producer.send(b'msg-0')))
        tasks.append(asyncio.create_task(producer.send(b'msg-1')))

        done, pending = await asyncio.wait(tasks, timeout=1, return_when=asyncio.FIRST_COMPLETED)
        self.assertEqual(len(done), 0)
        self.assertEqual(len(pending), 2)

        # flush will trigger sending the batched messages
        await producer.flush()
        for task in pending:
            self.assertTrue(task.done())
        msg_id0 = tasks[0].result()
        msg_id1 = tasks[1].result()
        self.assertEqual(msg_id0.ledger_id(), msg_id1.ledger_id())
        self.assertEqual(msg_id0.entry_id(), msg_id1.entry_id())
        self.assertEqual(msg_id0.batch_index(), 0)
        self.assertEqual(msg_id1.batch_index(), 1)

    async def test_create_producer_failure(self):
        try:
            await self._client.create_producer('tenant/ns/asyncio-test-send-failure')
            self.fail()
        except PulsarException as e:
            self.assertEqual(e.error(), pulsar.Result.Timeout)

    async def test_send_failure(self):
        producer = await self._client.create_producer('asyncio-test-send-failure')
        try:
            await producer.send(('x' * 1024 * 1024 * 10).encode())
            self.fail()
        except PulsarException as e:
            self.assertEqual(e.error(), pulsar.Result.MessageTooBig)

    async def test_close_producer(self):
        producer = await self._client.create_producer('asyncio-test-close-producer')
        await producer.close()
        try:
            await producer.close()
            self.fail()
        except PulsarException as e:
            self.assertEqual(e.error(), pulsar.Result.AlreadyClosed)

    async def _prepare_messages(self, producer: Producer) -> List[pulsar.MessageId]:
        msg_ids = []
        for i in range(5):
            msg_ids.append(await producer.send(f'msg-{i}'.encode()))
        return msg_ids

    async def test_consumer_cumulative_acknowledge(self):
        topic = f'asyncio-test-consumer-cumulative-ack-{time.time()}'
        sub = 'sub'
        consumer = await self._client.subscribe(topic, sub)
        producer = await self._client.create_producer(topic)
        await self._prepare_messages(producer)
        last_msg = None
        for _ in range(5):
            last_msg = await consumer.receive()
        await consumer.acknowledge_cumulative(last_msg)
        await consumer.close()

        consumer = await self._client.subscribe(topic, sub)
        await producer.send(b'final-message')
        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'final-message')

    async def test_consumer_individual_acknowledge(self):
        topic = f'asyncio-test-consumer-individual-ack-{time.time()}'
        sub = 'sub'
        consumer = await self._client.subscribe(topic, sub,
                                                consumer_type=pulsar.ConsumerType.Shared)
        producer = await self._client.create_producer(topic)
        await self._prepare_messages(producer)
        msgs = []
        for _ in range(5):
            msg = await consumer.receive()
            msgs.append(msg)

        await consumer.acknowledge(msgs[0])
        await consumer.acknowledge(msgs[2])
        await consumer.acknowledge(msgs[4])
        await consumer.close()

        consumer = await self._client.subscribe(topic, sub,
                                                consumer_type=pulsar.ConsumerType.Shared)
        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'msg-1')
        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'msg-3')

    async def test_multi_topic_consumer(self):
        topics = ['asyncio-test-multi-topic-1', 'asyncio-test-multi-topic-2']
        producers = []

        for topic in topics:
            producer = await self._client.create_producer(topic)
            producers.append(producer)

        consumer = await self._client.subscribe(topics, 'test-multi-subscription')

        await producers[0].send(b'message-from-topic-1')
        await producers[1].send(b'message-from-topic-2')

        async def verify_receive(consumer: Consumer):
            received_messages = {}
            for _ in range(2):
                msg = await consumer.receive()
                received_messages[msg.data()] = None
                await consumer.acknowledge(msg.message_id())
            self.assertEqual(received_messages, {
                b'message-from-topic-1': None,
                b'message-from-topic-2': None
            })

        await verify_receive(consumer)
        await consumer.close()

        consumer = await self._client.subscribe('public/default/asyncio-test-multi-topic-.*',
                                                'test-multi-subscription-2',
                                                is_pattern_topic=True,
                                                initial_position=pulsar.InitialPosition.Earliest)
        await verify_receive(consumer)
        await consumer.close()

    async def test_unsubscribe(self):
        topic = f'asyncio-test-unsubscribe-{time.time()}'
        sub = 'sub'
        consumer = await self._client.subscribe(topic, sub)
        await consumer.unsubscribe()
        # Verify the consumer can be created successfully with the same subscription name
        consumer = await self._client.subscribe(topic, sub)
        await consumer.close()

    async def test_seek_message_id(self):
        topic = f'asyncio-test-seek-message-id-{time.time()}'
        sub = 'sub'

        producer = await self._client.create_producer(topic)
        msg_ids = await self._prepare_messages(producer)

        consumer = await self._client.subscribe(
            topic, sub, initial_position=pulsar.InitialPosition.Earliest
        )
        await consumer.seek(msg_ids[2])
        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'msg-3')
        await consumer.close()

        consumer = await self._client.subscribe(
            topic, sub, initial_position=pulsar.InitialPosition.Earliest,
            start_message_id_inclusive=True
        )
        await consumer.seek(msg_ids[2])
        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'msg-2')
        await consumer.close()

    async def test_seek_timestamp(self):
        topic = f'asyncio-test-seek-timestamp-{time.time()}'
        sub = 'sub'
        consumer = await self._client.subscribe(
            topic, sub, initial_position=pulsar.InitialPosition.Earliest
        )

        producer = await self._client.create_producer(topic)

        # Send first 3 messages
        for i in range(3):
            await producer.send(f'msg-{i}'.encode())

        seek_time = int(time.time() * 1000)

        # Send 2 more messages
        for i in range(3, 5):
            await producer.send(f'msg-{i}'.encode())

        # Consume all messages first
        for i in range(5):
            msg = await consumer.receive()
            self.assertEqual(msg.data(), f'msg-{i}'.encode())

        # Seek to the timestamp (should start from msg-3)
        await consumer.seek(seek_time)

        msg = await consumer.receive()
        self.assertEqual(msg.data(), b'msg-3')

    async def test_schema(self):
        class ExampleRecord(Record):  # pylint: disable=too-few-public-methods
            """Example record schema for testing."""
            str_field = String()
            int_field = Integer()

        topic = f'asyncio-test-schema-{time.time()}'
        producer = await self._client.create_producer(
                topic, schema=AvroSchema(ExampleRecord)
        )
        consumer = await self._client.subscribe(
            topic, 'sub', schema=AvroSchema(ExampleRecord)
        )
        await producer.send(ExampleRecord(str_field='test', int_field=42))
        msg = await consumer.receive()
        self.assertIsInstance(msg.value(), ExampleRecord)
        self.assertEqual(msg.value().str_field, 'test')
        self.assertEqual(msg.value().int_field, 42)


if __name__ == '__main__':
    main()