From f8bb6bcc62232080cabe4824069f3cc8b877d9dd Mon Sep 17 00:00:00 2001
From: maxalbert <maximilian.albert@gmail.com>
Date: Fri, 21 Jun 2024 11:22:51 +0100
Subject: [PATCH] Add NatsBroker.new_inbox() (#1543)

* docs: update CONTRIBUTING.md

Added missing 'and not confluent' to pytest command for running the
tests when no broker instances are present.

* feat: enable spy_decorator to wrap both async and non-async methods

* feat: add NatsBroker.new_inbox()

* chore: linting

* fix: raise error if new_inbox() is called before broker was started

* refactor: use assert statement instead of raising an exception

* docs: mention 'inbox_prefix' argument in docstring for 'new_inbox()'
---
 .../contributing/CONTRIBUTING.md              |  2 +-
 faststream/nats/broker/broker.py              | 13 +++++++++++
 tests/brokers/nats/test_new_inbox.py          | 23 +++++++++++++++++++
 tests/brokers/nats/test_test_client.py        |  1 -
 tests/tools.py                                | 14 ++++++++---
 5 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 tests/brokers/nats/test_new_inbox.py

diff --git a/docs/docs/en/getting-started/contributing/CONTRIBUTING.md b/docs/docs/en/getting-started/contributing/CONTRIBUTING.md
index 9818aea06c..745f963830 100644
--- a/docs/docs/en/getting-started/contributing/CONTRIBUTING.md
+++ b/docs/docs/en/getting-started/contributing/CONTRIBUTING.md
@@ -94,7 +94,7 @@ pytest -m 'all'
 If you don't have a local broker instance running, you can run tests without those dependencies:
 
 ```bash
-pytest -m 'not rabbit and not kafka and not nats and not redis'
+pytest -m 'not rabbit and not kafka and not nats and not redis and not confluent'
 ```
 
 To run tests based on RabbitMQ, Kafka, or other dependencies, the following dependencies are needed to be started as docker containers:
diff --git a/faststream/nats/broker/broker.py b/faststream/nats/broker/broker.py
index 35e35086c8..fefcc069a8 100644
--- a/faststream/nats/broker/broker.py
+++ b/faststream/nats/broker/broker.py
@@ -901,3 +901,16 @@ async def wrapper() -> None:
                 self.__is_connected = True
 
         return wrapper
+
+    async def new_inbox(self) -> str:
+        """Return a unique inbox that can be used for NATS requests or subscriptions.
+
+        The inbox prefix can be customised by passing `inbox_prefix` when creating your `NatsBroker`.
+
+        This method calls `nats.aio.client.Client.new_inbox` [1] under the hood.
+
+        [1] https://nats-io.github.io/nats.py/modules.html#nats.aio.client.Client.new_inbox
+        """
+        assert self._connection  # nosec B101
+
+        return self._connection.new_inbox()
diff --git a/tests/brokers/nats/test_new_inbox.py b/tests/brokers/nats/test_new_inbox.py
new file mode 100644
index 0000000000..7969ad46a3
--- /dev/null
+++ b/tests/brokers/nats/test_new_inbox.py
@@ -0,0 +1,23 @@
+from unittest.mock import patch
+
+import pytest
+from nats.aio.client import Client as NatsClient
+
+from faststream.nats import NatsBroker
+from tests.tools import spy_decorator
+
+
+@pytest.mark.asyncio()
+@pytest.mark.nats()
+async def test_new_inbox():
+    with patch.object(
+        NatsClient,
+        "new_inbox",
+        spy_decorator(NatsClient.new_inbox),
+    ) as m:
+        broker = NatsBroker(inbox_prefix="_FOO_TEST_INBOX")
+        await broker.connect()
+        inbox_name = await broker.new_inbox()
+
+        m.mock.assert_called_once()
+        assert inbox_name.startswith("_FOO_TEST_INBOX.")
diff --git a/tests/brokers/nats/test_test_client.py b/tests/brokers/nats/test_test_client.py
index 94a4be5dac..e24cbb577e 100644
--- a/tests/brokers/nats/test_test_client.py
+++ b/tests/brokers/nats/test_test_client.py
@@ -91,7 +91,6 @@ async def test_inbox_prefix_with_real(
             assert br._connection._inbox_prefix == b"test"
             assert "test" in str(br._connection.new_inbox())
 
-
     async def test_respect_middleware(self, queue):
         routes = []
 
diff --git a/tests/tools.py b/tests/tools.py
index b1fc8e52c4..48df98a0a1 100644
--- a/tests/tools.py
+++ b/tests/tools.py
@@ -1,3 +1,4 @@
+import inspect
 from typing import Any, Iterable
 from unittest.mock import MagicMock
 
@@ -5,9 +6,16 @@
 def spy_decorator(method):
     mock = MagicMock()
 
-    async def wrapper(*args, **kwargs):
-        mock(*args, **kwargs)
-        return await method(*args, **kwargs)
+    if inspect.iscoroutinefunction(method):
+
+        async def wrapper(*args, **kwargs):
+            mock(*args, **kwargs)
+            return await method(*args, **kwargs)
+    else:
+
+        def wrapper(*args, **kwargs):
+            mock(*args, **kwargs)
+            return method(*args, **kwargs)
 
     wrapper.mock = mock
     return wrapper