feat: NATS test client supports filter subsciption

faststream/nats/subscriber/factory.py

@@ -4,6 +4,7 @@
     DEFAULT_SUB_PENDING_BYTES_LIMIT,
     DEFAULT_SUB_PENDING_MSGS_LIMIT,
 )
+from nats.js.api import ConsumerConfig
 from nats.js.client import (
     DEFAULT_JS_SUB_PENDING_BYTES_LIMIT,
     DEFAULT_JS_SUB_PENDING_MSGS_LIMIT,
@@ -83,6 +84,8 @@ def create_subscriber(
     if not subject and not config:
         raise SetupError("You must provide either `subject` or `config` option.")
 
+    config = config or ConsumerConfig(filter_subjects=[])
+
     if stream:
         # TODO: pull & queue warning
         # TODO: push & durable warning
@@ -94,7 +97,6 @@ def create_subscriber(
             or DEFAULT_JS_SUB_PENDING_BYTES_LIMIT,
             "durable": durable,
             "stream": stream.name,
-            "config": config,
         }
 
         if pull_sub is not None:
@@ -123,6 +125,7 @@ def create_subscriber(
     if obj_watch is not None:
         return AsyncAPIObjStoreWatchSubscriber(
             subject=subject,
+            config=config,
             obj_watch=obj_watch,
             broker_dependencies=broker_dependencies,
             broker_middlewares=broker_middlewares,
@@ -134,6 +137,7 @@ def create_subscriber(
     if kv_watch is not None:
         return AsyncAPIKeyValueWatchSubscriber(
             subject=subject,
+            config=config,
             kv_watch=kv_watch,
             broker_dependencies=broker_dependencies,
             broker_middlewares=broker_middlewares,
@@ -147,6 +151,7 @@ def create_subscriber(
             return AsyncAPIConcurrentCoreSubscriber(
                 max_workers=max_workers,
                 subject=subject,
+                config=config,
                 queue=queue,
                 # basic args
                 extra_options=extra_options,
@@ -165,6 +170,7 @@ def create_subscriber(
         else:
             return AsyncAPICoreSubscriber(
                 subject=subject,
+                config=config,
                 queue=queue,
                 # basic args
                 extra_options=extra_options,
@@ -188,6 +194,7 @@ def create_subscriber(
                     pull_sub=pull_sub,
                     stream=stream,
                     subject=subject,
+                    config=config,
                     # basic args
                     extra_options=extra_options,
                     # Subscriber args
@@ -207,6 +214,7 @@ def create_subscriber(
                     max_workers=max_workers,
                     stream=stream,
                     subject=subject,
+                    config=config,
                     queue=queue,
                     # basic args
                     extra_options=extra_options,
@@ -229,6 +237,7 @@ def create_subscriber(
                         pull_sub=pull_sub,
                         stream=stream,
                         subject=subject,
+                        config=config,
                         # basic args
                         extra_options=extra_options,
                         # Subscriber args
@@ -248,6 +257,7 @@ def create_subscriber(
                         pull_sub=pull_sub,
                         stream=stream,
                         subject=subject,
+                        config=config,
                         # basic args
                         extra_options=extra_options,
                         # Subscriber args
@@ -267,6 +277,7 @@ def create_subscriber(
                     stream=stream,
                     subject=subject,
                     queue=queue,
+                    config=config,
                     # basic args
                     extra_options=extra_options,
                     # Subscriber args

faststream/nats/subscriber/usecase.py

@@ -20,7 +20,7 @@
 import anyio
 from fast_depends.dependencies import Depends
 from nats.errors import ConnectionClosedError, TimeoutError
-from nats.js.api import ObjectInfo
+from nats.js.api import ConsumerConfig, ObjectInfo
 from nats.js.kv import KeyValue
 from typing_extensions import Annotated, Doc, override
 
@@ -73,6 +73,7 @@ def __init__(
         self,
         *,
         subject: str,
+        config: "ConsumerConfig",
         extra_options: Optional[AnyDict],
         # Subscriber args
         default_parser: "AsyncCallable",
@@ -88,6 +89,7 @@ def __init__(
         include_in_schema: bool,
     ) -> None:
         self.subject = subject
+        self.config = config
 
         self.extra_options = extra_options or {}
 
@@ -205,10 +207,18 @@ def build_log_context(
 
     def add_prefix(self, prefix: str) -> None:
         """Include Subscriber in router."""
-        self.subject = "".join((prefix, self.subject))
+        if self.subject:
+            self.subject = "".join((prefix, self.subject))
+        else:
+            self.config.filter_subjects = [
+                "".join((prefix, subject))
+                for subject in (self.config.filter_subjects or ())
+            ]
 
     def __hash__(self) -> int:
-        return self.get_routing_hash(self.subject)
+        return self.get_routing_hash(
+            self.subject or "".join(self.config.filter_subjects or ())
+        )
 
     @staticmethod
     def get_routing_hash(
@@ -229,6 +239,7 @@ def __init__(
         self,
         *,
         subject: str,
+        config: "ConsumerConfig",
         # default args
         extra_options: Optional[AnyDict],
         # Subscriber args
@@ -246,6 +257,7 @@ def __init__(
     ) -> None:
         super().__init__(
             subject=subject,
+            config=config,
             extra_options=extra_options,
             # subscriber args
             default_parser=default_parser,
@@ -368,6 +380,7 @@ def __init__(
         *,
         # default args
         subject: str,
+        config: "ConsumerConfig",
         queue: str,
         extra_options: Optional[AnyDict],
         # Subscriber args
@@ -387,6 +400,7 @@ def __init__(
 
         super().__init__(
             subject=subject,
+            config=config,
             extra_options=extra_options,
             # subscriber args
             default_parser=parser_.parse_message,
@@ -439,6 +453,7 @@ def __init__(
         max_workers: int,
         # default args
         subject: str,
+        config: "ConsumerConfig",
         queue: str,
         extra_options: Optional[AnyDict],
         # Subscriber args
@@ -456,6 +471,7 @@ def __init__(
             max_workers=max_workers,
             # basic args
             subject=subject,
+            config=config,
             queue=queue,
             extra_options=extra_options,
             # Propagated args
@@ -494,6 +510,7 @@ def __init__(
         stream: "JStream",
         # default args
         subject: str,
+        config: "ConsumerConfig",
         queue: str,
         extra_options: Optional[AnyDict],
         # Subscriber args
@@ -514,6 +531,7 @@ def __init__(
 
         super().__init__(
             subject=subject,
+            config=config,
             extra_options=extra_options,
             # subscriber args
             default_parser=parser_.parse_message,
@@ -540,7 +558,7 @@ def get_log_context(
         """Log context factory using in `self.consume` scope."""
         return self.build_log_context(
             message=message,
-            subject=self.subject,
+            subject=self.subject or ", ".join(self.config.filter_subjects or ()),
             queue=self.queue,
             stream=self.stream.name,
         )
@@ -560,6 +578,7 @@ async def _create_subscription(  # type: ignore[override]
             subject=self.clear_subject,
             queue=self.queue,
             cb=self.consume,
+            config=self.config,
             **self.extra_options,
         )
 
@@ -574,6 +593,7 @@ def __init__(
         stream: "JStream",
         # default args
         subject: str,
+        config: "ConsumerConfig",
         queue: str,
         extra_options: Optional[AnyDict],
         # Subscriber args
@@ -592,6 +612,7 @@ def __init__(
             # basic args
             stream=stream,
             subject=subject,
+            config=config,
             queue=queue,
             extra_options=extra_options,
             # Propagated args
@@ -619,6 +640,7 @@ async def _create_subscription(  # type: ignore[override]
             subject=self.clear_subject,
             queue=self.queue,
             cb=self._put_msg,
+            config=self.config,
             **self.extra_options,
         )
 
@@ -633,6 +655,7 @@ def __init__(
         stream: "JStream",
         # default args
         subject: str,
+        config: "ConsumerConfig",
         extra_options: Optional[AnyDict],
         # Subscriber args
         no_ack: bool,
@@ -651,6 +674,7 @@ def __init__(
             # basic args
             stream=stream,
             subject=subject,
+            config=config,
             extra_options=extra_options,
             queue="",
             # Propagated args
@@ -708,6 +732,7 @@ def __init__(
         pull_sub: "PullSub",
         stream: "JStream",
         subject: str,
+        config: "ConsumerConfig",
         extra_options: Optional[AnyDict],
         # Subscriber args
         no_ack: bool,
@@ -726,6 +751,7 @@ def __init__(
             pull_sub=pull_sub,
             stream=stream,
             subject=subject,
+            config=config,
             extra_options=extra_options,
             # Propagated args
             no_ack=no_ack,
@@ -765,6 +791,7 @@ def __init__(
         *,
         # default args
         subject: str,
+        config: "ConsumerConfig",
         stream: "JStream",
         pull_sub: "PullSub",
         extra_options: Optional[AnyDict],
@@ -786,6 +813,7 @@ def __init__(
 
         super().__init__(
             subject=subject,
+            config=config,
             extra_options=extra_options,
             # subscriber args
             default_parser=parser.parse_batch,
@@ -837,6 +865,7 @@ def __init__(
         self,
         *,
         subject: str,
+        config: "ConsumerConfig",
         kv_watch: "KvWatch",
         broker_dependencies: Iterable[Depends],
         broker_middlewares: Iterable["BrokerMiddleware[KeyValue.Entry]"],
@@ -850,6 +879,7 @@ def __init__(
 
         super().__init__(
             subject=subject,
+            config=config,
             extra_options=None,
             no_ack=True,
             no_reply=True,
@@ -941,6 +971,7 @@ def __init__(
         self,
         *,
         subject: str,
+        config: "ConsumerConfig",
         obj_watch: "ObjWatch",
         broker_dependencies: Iterable[Depends],
         broker_middlewares: Iterable["BrokerMiddleware[List[Msg]]"],
@@ -955,6 +986,7 @@ def __init__(
 
         super().__init__(
             subject=subject,
+            config=config,
             extra_options=None,
             no_ack=True,
             no_reply=True,

faststream/nats/testing.py

@@ -97,7 +97,10 @@ async def publish(  # type: ignore[override]
             ):
                 continue
 
-            if is_subject_match_wildcard(subject, handler.clear_subject):
+            if is_subject_match_wildcard(subject, handler.clear_subject) or any(
+                is_subject_match_wildcard(subject, filter_subject)
+                for filter_subject in (handler.config.filter_subjects or ())
+            ):
                 msg: Union[List[PatchedMessage], PatchedMessage]
                 if (pull := getattr(handler, "pull_sub", None)) and pull.batch:
                     msg = [incoming]

tests/brokers/nats/test_consume.py

@@ -1,11 +1,11 @@
 import asyncio
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 from nats.aio.msg import Msg
 
 from faststream.exceptions import AckMessage
-from faststream.nats import JStream, NatsBroker, PullSub
+from faststream.nats import ConsumerConfig, JStream, NatsBroker, PullSub
 from faststream.nats.annotations import NatsMessage
 from tests.brokers.base.consume import BrokerRealConsumeTestcase
 from tests.tools import spy_decorator
@@ -40,6 +40,36 @@ def subscriber(m):
 
         assert event.is_set()
 
+    async def test_consume_with_filter(
+        self,
+        queue,
+        mock: Mock,
+        event: asyncio.Event,
+    ):
+        consume_broker = self.get_broker()
+
+        @consume_broker.subscriber(
+            config=ConsumerConfig(filter_subjects=[f"{queue}.a"]),
+            stream=JStream(queue, subjects=[f"{queue}.*"]),
+        )
+        def subscriber(m):
+            mock(m)
+            event.set()
+
+        async with self.patch_broker(consume_broker) as br:
+            await br.start()
+            await asyncio.wait(
+                (
+                    asyncio.create_task(br.publish(1, f"{queue}.b")),
+                    asyncio.create_task(br.publish(2, f"{queue}.a")),
+                    asyncio.create_task(event.wait()),
+                ),
+                timeout=3,
+            )
+
+        assert event.is_set()
+        mock.assert_called_once_with(2)
+
     async def test_consume_pull(
         self,
         queue: str,