Added tests for async client

Author: nl-x-net

pulsar/asyncio.py

@@ -23,11 +23,12 @@
 
 import asyncio
 import functools
-from typing import Any, Iterable, Optional
+from typing import Any, Iterable, Optional, Union
 
 import _pulsar
 
 import pulsar
+from pulsar import _listener_wrapper
 
 
 class PulsarException(BaseException):
@@ -212,12 +213,16 @@ async def close(self):
         self._consumer.close_async(functools.partial(_set_future, future, value=None))
         await future
 
-    async def seek(self, position: tuple[int, int, int, int]):
+    async def seek(self, position: tuple[int, int, int, int] | pulsar.MessageId):
         """
         Reset the subscription associated with this consumer to a specific message id or publish timestamp. The message id can either be a specific message or represent the first or last messages in the topic. ...
         """
-        partition, ledger_id, entry_id, batch_index = position
-        message_id = _pulsar.MessageId(partition, ledger_id, entry_id, batch_index)
+        if isinstance(position, tuple):
+            partition, ledger_id, entry_id, batch_index = position
+            message_id = _pulsar.MessageId(partition, ledger_id, entry_id, batch_index)
+        else:
+            message_id = position
+
         future = asyncio.get_running_loop().create_future()
         self._consumer.seek_async(message_id, functools.partial(_set_future, future))
         await future
@@ -249,21 +254,40 @@ def redeliver_unacknowledged_messages(self):
         self._consumer.redeliver_unacknowledged_messages()
 
     @property
-    def last_message_id(self) -> int:
-        return self._consumer.last_message_id
+    def last_message_id(self) -> pulsar.MessageId:
+        """
+        MessageId of the last consumed message
+        """
+        return self._consumer.get_last_message_id()
 
     @property
     def is_connected(self) -> bool:
+        """
+        True if the consumer is connected to a broker
+        """
         return self._consumer.is_connected()
 
     @property
     def subscription_name(self) -> str:
+        """
+        Name of the current subscription
+        """
         return self._consumer.subscription_name()
 
     @property
     def topic(self) -> str:
+        """
+        Topic(s) of consumer
+        """
         return self._consumer.topic()
 
+    @property
+    def consumer_name(self) -> str:
+        """
+        Name of consumer
+        """
+        return self._consumer.consumer_name()
+
 
 class Client:
     """
@@ -277,17 +301,104 @@ def __init__(self, service_url, **kwargs) -> None:
         assert service_url.startswith('pulsar://'), "The service url must start with 'pulsar://'"
         self._client = pulsar.Client(service_url, **kwargs)._client
 
-    async def subscribe(self, topics: str, subscription_name: str, consumer_type: _pulsar.ConsumerType,
-                        schema: Optional[_pulsar.SchemaInfo] = _pulsar.SchemaInfo(_pulsar.SchemaType.BYTES, "bytes", "")) -> Consumer:
+    async def subscribe(self, topic, subscription_name,
+                        consumer_type: _pulsar.ConsumerType = _pulsar.ConsumerType.Exclusive,
+                        schema=pulsar.schema.BytesSchema(),
+                        message_listener=None,
+                        receiver_queue_size=1000,
+                        max_total_receiver_queue_size_across_partitions=50000,
+                        consumer_name=None,
+                        unacked_messages_timeout_ms=None,
+                        broker_consumer_stats_cache_time_ms=30000,
+                        negative_ack_redelivery_delay_ms=60000,
+                        is_read_compacted=False,
+                        properties=None,
+                        pattern_auto_discovery_period=60,
+                        initial_position: _pulsar.InitialPosition = _pulsar.InitialPosition.Latest,
+                        crypto_key_reader: Union[None, _pulsar.CryptoKeyReader] = None,
+                        replicate_subscription_state_enabled=False,
+                        max_pending_chunked_message=10,
+                        auto_ack_oldest_chunked_message_on_queue_full=False,
+                        start_message_id_inclusive=False,
+                        batch_receive_policy=None,
+                        key_shared_policy=None,
+                        batch_index_ack_enabled=False,
+                        regex_subscription_mode: _pulsar.RegexSubscriptionMode = _pulsar.RegexSubscriptionMode.PersistentOnly,
+                        dead_letter_policy: Union[None, pulsar.ConsumerDeadLetterPolicy] = None,) -> Consumer:
         conf = _pulsar.ConsumerConfiguration()
         conf.consumer_type(consumer_type)
-        conf.schema(schema)
+        conf.regex_subscription_mode(regex_subscription_mode)
+        conf.read_compacted(is_read_compacted)
+        if message_listener:
+            conf.message_listener(_listener_wrapper(message_listener, schema))
+        conf.receiver_queue_size(receiver_queue_size)
+        conf.max_total_receiver_queue_size_across_partitions(max_total_receiver_queue_size_across_partitions)
+        if consumer_name:
+            conf.consumer_name(consumer_name)
+        if unacked_messages_timeout_ms:
+            conf.unacked_messages_timeout_ms(unacked_messages_timeout_ms)
+
+        conf.negative_ack_redelivery_delay_ms(negative_ack_redelivery_delay_ms)
+        conf.broker_consumer_stats_cache_time_ms(broker_consumer_stats_cache_time_ms)
+        if properties:
+            for k, v in properties.items():
+                conf.property(k, v)
+        conf.subscription_initial_position(initial_position)
+
+        conf.schema(schema.schema_info())
+
+        if crypto_key_reader:
+            conf.crypto_key_reader(crypto_key_reader.cryptoKeyReader)
+
+        conf.replicate_subscription_state_enabled(replicate_subscription_state_enabled)
+        conf.max_pending_chunked_message(max_pending_chunked_message)
+        conf.auto_ack_oldest_chunked_message_on_queue_full(auto_ack_oldest_chunked_message_on_queue_full)
+        conf.start_message_id_inclusive(start_message_id_inclusive)
+        if batch_receive_policy:
+            conf.batch_receive_policy(batch_receive_policy.policy())
+
+        if key_shared_policy:
+            conf.key_shared_policy(key_shared_policy.policy())
+        conf.batch_index_ack_enabled(batch_index_ack_enabled)
+        if dead_letter_policy:
+            conf.dead_letter_policy(dead_letter_policy.policy())
 
         future = asyncio.get_running_loop().create_future()
-        self._client.subscribe_async(topics, subscription_name, conf, functools.partial(_set_future, future))
+
+        if isinstance(topic, str):
+            self._client.subscribe_async(topic, subscription_name, conf, functools.partial(_set_future, future))
+        elif isinstance(topic, list):
+            self._client.subscribe_topics_async(topic, subscription_name, conf, functools.partial(_set_future, future))
+        elif isinstance(topic, pulsar._retype):
+            self._client.subscribe_pattern_async(topic, subscription_name, conf, functools.partial(_set_future, future))
+        else:
+            raise ValueError("Argument 'topic' is expected to be of a type between (str, list, re.pattern)")
+
         return Consumer(await future)
 
-    async def create_producer(self, topic: str) -> Producer:
+    async def create_producer(self, topic,
+                              producer_name=None,
+                              schema=pulsar.schema.BytesSchema(),
+                              initial_sequence_id=None,
+                              send_timeout_millis=30000,
+                              compression_type: _pulsar.CompressionType = _pulsar.CompressionType.NONE,
+                              max_pending_messages=1000,
+                              max_pending_messages_across_partitions=50000,
+                              block_if_queue_full=False,
+                              batching_enabled=False,
+                              batching_max_messages=1000,
+                              batching_max_allowed_size_in_bytes=128*1024,
+                              batching_max_publish_delay_ms=10,
+                              chunking_enabled=False,
+                              message_routing_mode: _pulsar.PartitionsRoutingMode = _pulsar.PartitionsRoutingMode.RoundRobinDistribution,
+                              lazy_start_partitioned_producers=False,
+                              properties=None,
+                              batching_type: _pulsar.BatchingType = _pulsar.BatchingType.Default,
+                              encryption_key=None,
+                              crypto_key_reader: Union[None, _pulsar.CryptoKeyReader] = None,
+                              access_mode: _pulsar.ProducerAccessMode = _pulsar.ProducerAccessMode.Shared,
+
+                              ) -> Producer:
         """
         Create a new producer on a given topic
 
@@ -306,8 +417,39 @@ async def create_producer(self, topic: str) -> Producer:
         PulsarException
         """
         future = asyncio.get_running_loop().create_future()
+
         conf = _pulsar.ProducerConfiguration()
-        # TODO: add more configs
+        conf.send_timeout_millis(send_timeout_millis)
+        conf.compression_type(compression_type)
+        conf.max_pending_messages(max_pending_messages)
+        conf.max_pending_messages_across_partitions(max_pending_messages_across_partitions)
+        conf.block_if_queue_full(block_if_queue_full)
+        conf.batching_enabled(batching_enabled)
+        conf.batching_max_messages(batching_max_messages)
+        conf.batching_max_allowed_size_in_bytes(batching_max_allowed_size_in_bytes)
+        conf.batching_max_publish_delay_ms(batching_max_publish_delay_ms)
+        conf.partitions_routing_mode(message_routing_mode)
+        conf.batching_type(batching_type)
+        conf.chunking_enabled(chunking_enabled)
+        conf.lazy_start_partitioned_producers(lazy_start_partitioned_producers)
+        conf.access_mode(access_mode)
+        if producer_name:
+            conf.producer_name(producer_name)
+        if initial_sequence_id:
+            conf.initial_sequence_id(initial_sequence_id)
+        if properties:
+            for k, v in properties.items():
+                conf.property(k, v)
+
+        conf.schema(schema.schema_info())
+        if encryption_key:
+            conf.encryption_key(encryption_key)
+        if crypto_key_reader:
+            conf.crypto_key_reader(crypto_key_reader.cryptoKeyReader)
+
+        if batching_enabled and chunking_enabled:
+            raise ValueError("Batching and chunking can´t be enabled at the same time")
+
         self._client.create_producer_async(topic, conf, functools.partial(_set_future, future))
         return Producer(await future)
 
@@ -340,3 +482,4 @@ def complete():
             future.set_exception(PulsarException(result))
 
     future.get_loop().call_soon_threadsafe(complete)
+

src/client.cc

@@ -21,6 +21,7 @@
 #include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <cmath>
 
 namespace py = pybind11;
 
@@ -53,13 +54,27 @@ Consumer Client_subscribe_topics(Client& client, const std::vector<std::string>&
         [&](SubscribeCallback callback) { client.subscribeAsync(topics, subscriptionName, conf, callback); });
 }
 
+void Client_subscribe_topicsAsync(Client& client, const std::vector<std::string>& topics, const std::string& subscriptionName, const ConsumerConfiguration& conf, SubscribeCallback callback){
+    client.subscribeAsync(topics, subscriptionName, conf, [callback](Result result, pulsar::Consumer consumer){
+        py::gil_scoped_acquire acquire;
+        callback(result, consumer);
+    });
+}
+
 Consumer Client_subscribe_pattern(Client& client, const std::string& topic_pattern,
                                   const std::string& subscriptionName, const ConsumerConfiguration& conf) {
     return waitForAsyncValue<Consumer>([&](SubscribeCallback callback) {
         client.subscribeWithRegexAsync(topic_pattern, subscriptionName, conf, callback);
     });
 }
 
+void Client_subscribe_patternAsync(Client& client, const std::string& topic_pattern, const std::string& subscriptionName, const ConsumerConfiguration& conf, SubscribeCallback callback){
+    client.subscribeWithRegexAsync(topic_pattern, subscriptionName, conf, [callback](Result result, Consumer consumer){
+        py::gil_scoped_acquire acquire;
+        callback(result, consumer);
+    });
+}
+
 Reader Client_createReader(Client& client, const std::string& topic, const MessageId& startMessageId,
                            const ReaderConfiguration& conf) {
     return waitForAsyncValue<Reader>(
@@ -94,7 +109,9 @@ void export_client(py::module_& m) {
         .def("subscribe", &Client_subscribe)
         .def("subscribe_async", &Client_subscribeAsync)
         .def("subscribe_topics", &Client_subscribe_topics)
+        .def("subscribe_topics_async", &Client_subscribe_topicsAsync)
         .def("subscribe_pattern", &Client_subscribe_pattern)
+        .def("subscribe_pattern_async", &Client_subscribe_patternAsync)
         .def("create_reader", &Client_createReader)
         .def("get_topic_partitions", &Client_getTopicPartitions)
         .def("get_schema_info", &Client_getSchemaInfo)

tests/asyncio_test.py

@@ -19,13 +19,15 @@
 #
 
 import asyncio
+from typing import Iterable
 
 from _pulsar import ConsumerType
 
 import pulsar
 from pulsar.asyncio import (
     Client,
     PulsarException,
+    Consumer
 )
 from unittest import (
     main,
@@ -59,8 +61,6 @@ async def test_batch_send(self):
             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)
 
     async def test_create_producer_failure(self):
         try:
@@ -87,5 +87,53 @@ async def test_close_producer(self):
         except PulsarException as e:
             self.assertEqual(e.error(), pulsar.Result.AlreadyClosed)
 
+    async def test_subscribe(self):
+        consumer = await self._client.subscribe('awaitio-test-close-producer', 'test-subscription')
+        self.assertIsInstance(consumer, Consumer)
+
+    async def test_read_and_ack(self):
+        test_producer = await self._client.create_producer("awaitio-test-consumer-ack")
+        consumer = await self._client.subscribe('awaitio-test-consumer-ack', 'test-subscription')
+
+        await test_producer.send(b"test123")
+        msg = await consumer.receive()
+
+        self.assertEqual(msg.data(), b"test123")
+
+        await consumer.acknowledge(msg)
+
+    async def test_batch_read_and_ack(self):
+        test_producer = await self._client.create_producer("awaitio-test-consumer-ack-batch")
+        consumer = await self._client.subscribe('awaitio-test-consumer-ack-batch', 'test-subscription')
+
+        await test_producer.send(b"test123")
+        msgs = await consumer.batch_receive()
+
+        last = None
+        for msg in msgs:
+            last = msg
+
+        await consumer.acknowledge_cumulative(last)
+
+        self.assertIsInstance(msgs, Iterable)
+        for msg in msgs:
+            self.assertEqual(b"test123", msg.data())
+
+    async def test_consumer_close(self):
+        consumer = await self._client.subscribe('awaitio-test-consumer-close', 'test-subscription')
+        await consumer.close()
+
+        self.assertFalse(consumer.is_connected)
+
+    async def test_consumer_seek(self):
+        consumer = await self._client.subscribe('awaitio-test-consumer-close', 'test-subscription')
+        await consumer.seek(consumer.last_message_id)
+
+    async def test_consumer_unsubscribe(self):
+        consumer = await self._client.subscribe('awaitio-test-consumer-close', 'test-subscription')
+        await consumer.unsubscribe()
+
+    
+
 if __name__ == '__main__':
     main()