Fixed tests, and jetstream publish method.

Author: s3rius

python/natsrpy/_natsrpy_rs/js/__init__.pyi

@@ -1,18 +1,46 @@
 from datetime import datetime, timedelta
-from typing import Any
+from typing import Any, Literal, overload
 
 from .managers import KVManager, ObjectStoreManager, StreamsManager
 
+class Publication:
+    stream: str
+    sequence: int
+    domain: str
+    duplicate: bool
+    value: str | None
+
 class JetStream:
+    @overload
+    async def publish(
+        self,
+        subject: str,
+        payload: str | bytes | bytearray | memoryview,
+        *,
+        headers: dict[str, str] | None = None,
+        err_on_disconnect: bool = False,
+        wait: Literal[True],
+    ) -> Publication: ...
+    @overload
     async def publish(
         self,
         subject: str,
         payload: str | bytes | bytearray | memoryview,
         *,
         headers: dict[str, str] | None = None,
-        reply: str | None = None,
         err_on_disconnect: bool = False,
+        wait: Literal[False] = False,
     ) -> None: ...
+    @overload
+    async def publish(
+        self,
+        subject: str,
+        payload: str | bytes | bytearray | memoryview,
+        *,
+        headers: dict[str, str] | None = None,
+        err_on_disconnect: bool = False,
+        wait: bool = False,
+    ) -> Publication | None: ...
     @property
     def kv(self) -> KVManager: ...
     @property

python/natsrpy/_natsrpy_rs/js/stream.pyi

@@ -228,15 +228,22 @@ class StreamInfo:
     sources: list[SourceInfo]
 
 class Stream:
-    async def direct_get(self, sequence: int) -> StreamMessage:
+    async def direct_get(
+        self,
+        sequence: int,
+        timeout: float | datetime | None = None,
+    ) -> StreamMessage:
         """
         Get direct message from the stream.
 
         :param sequence: sequence number of the message to get.
         :return: Message.
         """
 
-    async def get_info(self) -> StreamInfo:
+    async def get_info(
+        self,
+        timeout: float | datetime | None = None,
+    ) -> StreamInfo:
         """
         Get information about the stream.
 
@@ -248,6 +255,7 @@ class Stream:
         filter: str | None = None,
         sequence: int | None = None,
         keep: int | None = None,
+        timeout: float | datetime | None = None,
     ) -> int:
         """
         Purge current stream.

python/tests/conftest.py

@@ -3,6 +3,7 @@
 
 import pytest
 from natsrpy import Nats
+from natsrpy.js import JetStream
 
 
 @pytest.fixture(scope="session")
@@ -29,3 +30,8 @@ async def nats(nats_url: str) -> AsyncGenerator[Nats, None]:
     yield nats
 
     await nats.shutdown()
+
+
+@pytest.fixture(scope="session")
+async def js(nats: Nats) -> JetStream:
+    return await nats.jetstream()

python/tests/test_consumers.py

@@ -1,7 +1,5 @@
 import uuid
 
-import pytest
-from natsrpy import Nats
 from natsrpy.js import (
     AckPolicy,
     DeliverPolicy,
@@ -15,11 +13,6 @@
 )
 
 
-@pytest.fixture()
-async def js(nats: Nats) -> JetStream:
-    return await nats.jetstream()
-
-
 async def test_pull_consumer_create(js: JetStream) -> None:
     stream_name = f"test-pcreate-{uuid.uuid4()}"
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
@@ -40,7 +33,7 @@ async def test_pull_consumer_fetch_with_ack(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"ack-msg")
+        await js.publish(subj, b"ack-msg", wait=True)
 
         consumer_config = PullConsumerConfig(
             name=f"consumer-{uuid.uuid4()}",
@@ -60,7 +53,7 @@ async def test_pull_consumer_nack(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"nack-msg")
+        await js.publish(subj, b"nack-msg", wait=True)
 
         consumer_config = PullConsumerConfig(
             name=f"consumer-{uuid.uuid4()}",
@@ -80,7 +73,7 @@ async def test_pull_consumer_term(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"term-msg")
+        await js.publish(subj, b"term-msg", wait=True)
 
         consumer_config = PullConsumerConfig(
             name=f"consumer-{uuid.uuid4()}",
@@ -100,7 +93,7 @@ async def test_pull_consumer_progress(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"progress-msg")
+        await js.publish(subj, b"progress-msg", wait=True)
 
         consumer_config = PullConsumerConfig(
             name=f"consumer-{uuid.uuid4()}",
@@ -121,7 +114,7 @@ async def test_pull_consumer_message_properties(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"prop-msg")
+        await js.publish(subj, b"prop-msg", wait=True)
 
         consumer_name = f"consumer-{uuid.uuid4()}"
         consumer_config = PullConsumerConfig(name=consumer_name)
@@ -149,8 +142,8 @@ async def test_pull_consumer_with_filter_subject(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(f"{stream_name}.a", b"msg-a")
-        await js.publish(f"{stream_name}.b", b"msg-b")
+        await js.publish(f"{stream_name}.a", b"msg-a", wait=True)
+        await js.publish(f"{stream_name}.b", b"msg-b", wait=True)
 
         consumer_config = PullConsumerConfig(
             name=f"consumer-{uuid.uuid4()}",
@@ -171,8 +164,8 @@ async def test_pull_consumer_deliver_policy(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"old-msg")
-        await js.publish(subj, b"new-msg")
+        await js.publish(subj, b"old-msg", wait=True)
+        await js.publish(subj, b"new-msg", wait=True)
 
         consumer_config = PullConsumerConfig(
             name=f"consumer-{uuid.uuid4()}",
@@ -243,7 +236,7 @@ async def test_pull_consumer_messages(js: JetStream) -> None:
     stream = await js.streams.create(config)
     try:
         for message in messages:
-            await js.publish(subj, message)
+            await js.publish(subj, message, wait=True)
         consumer_config = PullConsumerConfig(name=f"consumer-{uuid.uuid4()}")
         consumer = await stream.consumers.create(consumer_config)
         msgs_iter = await consumer.fetch(timeout=0.5)
@@ -261,7 +254,7 @@ async def test_push_consumer_messages(js: JetStream) -> None:
     stream = await js.streams.create(config)
     try:
         for message in messages:
-            await js.publish(subj, message)
+            await js.publish(subj, message, wait=True)
 
         deliver_subj = uuid.uuid4().hex
         consumer_config = PushConsumerConfig(

python/tests/test_jetstream.py

@@ -8,11 +8,6 @@
 )
 
 
-@pytest.fixture()
-async def js(nats: Nats) -> JetStream:
-    return await nats.jetstream()
-
-
 async def test_jetstream_creation(nats: Nats) -> None:
     js = await nats.jetstream()
     assert isinstance(js, JetStream)
@@ -30,26 +25,14 @@ async def test_jetstream_has_object_store_manager(js: JetStream) -> None:
     assert js.object_store is not None
 
 
-async def test_jetstream_publish(js: JetStream) -> None:
-    stream_name = f"test-js-pub-{uuid.uuid4().hex[:8]}"
-    subj = f"{stream_name}.data"
-    config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
-    stream = await js.streams.create(config)
-    try:
-        await js.publish(subj, b"jetstream-msg")
-        info = await stream.get_info()
-        assert info.state.messages >= 1
-    finally:
-        await js.streams.delete(stream_name)
-
-
-async def test_jetstream_publish_str(js: JetStream) -> None:
-    stream_name = f"test-js-pubstr-{uuid.uuid4().hex[:8]}"
+@pytest.mark.parametrize("payload", [b"bytes-test", "str-test"])
+async def test_jetstream_publish(js: JetStream, payload: str | bytes) -> None:
+    stream_name = f"test-js-pub-{uuid.uuid4().hex}"
     subj = f"{stream_name}.data"
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, "string-payload")
+        await js.publish(subj, payload, wait=True)
         info = await stream.get_info()
         assert info.state.messages >= 1
     finally:
@@ -62,6 +45,6 @@ async def test_jetstream_publish_with_headers(js: JetStream) -> None:
     config = StreamConfig(name=stream_name, subjects=[f"{stream_name}.>"])
     await js.streams.create(config)
     try:
-        await js.publish(subj, b"with-headers", headers={"x-test": "value"})
+        await js.publish(subj, b"with-headers", headers={"x-test": "value"}, wait=True)
     finally:
         await js.streams.delete(stream_name)

python/tests/test_kv.py

@@ -1,7 +1,5 @@
 import uuid
 
-import pytest
-from natsrpy import Nats
 from natsrpy.js import (
     JetStream,
     KeyValue,
@@ -10,11 +8,6 @@
 )
 
 
-@pytest.fixture()
-async def js(nats: Nats) -> JetStream:
-    return await nats.jetstream()
-
-
 async def test_kv_create(js: JetStream) -> None:
     bucket = f"test-kv-create-{uuid.uuid4().hex[:8]}"
     config = KVConfig(bucket=bucket)

python/tests/test_object_store.py

@@ -1,8 +1,6 @@
 import io
 import uuid
 
-import pytest
-from natsrpy import Nats
 from natsrpy.js import (
     JetStream,
     ObjectStore,
@@ -11,11 +9,6 @@
 )
 
 
-@pytest.fixture()
-async def js(nats: Nats) -> JetStream:
-    return await nats.jetstream()
-
-
 async def test_object_store_create(js: JetStream) -> None:
     bucket = f"test-os-create-{uuid.uuid4().hex[:8]}"
     config = ObjectStoreConfig(bucket=bucket)

python/tests/test_streams.py

@@ -1,7 +1,5 @@
 import uuid
 
-import pytest
-from natsrpy import Nats
 from natsrpy.js import (
     DiscardPolicy,
     JetStream,
@@ -12,11 +10,6 @@
 )
 
 
-@pytest.fixture()
-async def js(nats: Nats) -> JetStream:
-    return await nats.jetstream()
-
-
 async def test_stream_create(js: JetStream) -> None:
     name = f"test-create-{uuid.uuid4().hex[:8]}"
     config = StreamConfig(name=name, subjects=[f"{name}.>"])
@@ -94,9 +87,9 @@ async def test_stream_purge(js: JetStream) -> None:
     config = StreamConfig(name=name, subjects=[f"{name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"msg-1")
-        await js.publish(subj, b"msg-2")
-        await js.publish(subj, b"msg-3")
+        await js.publish(subj, b"msg-1", wait=True)
+        await js.publish(subj, b"msg-2", wait=True)
+        await js.publish(subj, b"msg-3", wait=True)
         info = await stream.get_info()
         assert info.state.messages == 3
         purged = await stream.purge()
@@ -112,8 +105,8 @@ async def test_stream_purge_with_filter(js: JetStream) -> None:
     config = StreamConfig(name=name, subjects=[f"{name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(f"{name}.a", b"a-msg")
-        await js.publish(f"{name}.b", b"b-msg")
+        await js.publish(f"{name}.a", b"a-msg", wait=True)
+        await js.publish(f"{name}.b", b"b-msg", wait=True)
         purged = await stream.purge(filter=f"{name}.a")
         assert purged == 1
         info = await stream.get_info()
@@ -128,7 +121,7 @@ async def test_stream_direct_get(js: JetStream) -> None:
     config = StreamConfig(name=name, subjects=[f"{name}.>"], allow_direct=True)
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"direct-get-msg")
+        await js.publish(subj, b"direct-get-msg", wait=True)
         msg = await stream.direct_get(sequence=1)
         assert msg.payload == b"direct-get-msg"
         assert msg.subject == subj
@@ -143,7 +136,7 @@ async def test_stream_message_repr(js: JetStream) -> None:
     config = StreamConfig(name=name, subjects=[f"{name}.>"], allow_direct=True)
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"repr-test")
+        await js.publish(subj, b"repr-test", wait=True)
         msg = await stream.direct_get(sequence=1)
         r = repr(msg)
         assert isinstance(r, str)
@@ -250,8 +243,8 @@ async def test_stream_state_after_publish(js: JetStream) -> None:
     config = StreamConfig(name=name, subjects=[f"{name}.>"])
     stream = await js.streams.create(config)
     try:
-        await js.publish(subj, b"msg-1")
-        await js.publish(subj, b"msg-2")
+        await js.publish(subj, b"msg-1", wait=True)
+        await js.publish(subj, b"msg-2", wait=True)
         info = await stream.get_info()
         assert info.state.messages == 2
         assert info.state.first_sequence == 1

src/exceptions/rust_err.rs

@@ -47,17 +47,19 @@ pub enum NatsrpyError {
     #[error(transparent)]
     CreateKeyValueError(#[from] async_nats::jetstream::context::CreateKeyValueError),
     #[error(transparent)]
-    KVEntryError(#[from] async_nats::jetstream::kv::EntryError),
+    CreateStreamError(#[from] async_nats::jetstream::context::CreateStreamError),
     #[error(transparent)]
-    KVPutError(#[from] async_nats::jetstream::kv::PutError),
+    GetStreamError(#[from] async_nats::jetstream::context::GetStreamError),
     #[error(transparent)]
     KVUpdateError(#[from] async_nats::jetstream::context::UpdateKeyValueError),
     #[error(transparent)]
-    DeleteError(#[from] async_nats::jetstream::kv::DeleteError),
+    JSPublishError(#[from] async_nats::jetstream::context::PublishError),
     #[error(transparent)]
-    CreateStreamError(#[from] async_nats::jetstream::context::CreateStreamError),
+    KVEntryError(#[from] async_nats::jetstream::kv::EntryError),
     #[error(transparent)]
-    GetStreamError(#[from] async_nats::jetstream::context::GetStreamError),
+    KVPutError(#[from] async_nats::jetstream::kv::PutError),
+    #[error(transparent)]
+    DeleteError(#[from] async_nats::jetstream::kv::DeleteError),
     #[error(transparent)]
     StreamDirectGetError(#[from] async_nats::jetstream::stream::DirectGetError),
     #[error(transparent)]