Feature: emulate Redis Stream PEL in TestRedisBroker to route messages to min_idle_time subscribers

faststream/redis/testing.py

@@ -1,7 +1,9 @@
 import re
+import uuid
 from collections.abc import AsyncGenerator, Iterable, Iterator, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack, asynccontextmanager, contextmanager
+from dataclasses import dataclass
 from functools import partial
 from typing import (
     TYPE_CHECKING,
@@ -56,6 +58,26 @@
 __all__ = ("TestRedisBroker",)
 
 
+@dataclass(kw_only=True)
+class Entry:
+    handler: Any
+    msg: Any
+
+
+class PEL:
+    def __init__(self) -> None:
+        self._entries: dict[str, Entry] = {}
+
+    def remove(self, correlation_id: Any) -> None:
+        self._entries.pop(correlation_id)
+
+    def put(self, msg: Any, handler: Any, correlation_id: Any) -> None:
+        self._entries.update({correlation_id: Entry(msg=msg, handler=handler)})
+
+    def get_entry(self, correlation_id: Any) -> Entry | None:
+        return self._entries.get(correlation_id)
+
+
 class TestRedisBroker(TestBroker[RedisBroker, EnterType]):
     """A class to test Redis brokers."""
 
@@ -82,7 +104,11 @@ def __init__(
         *brokers: RedisBroker,
         with_real: bool = False,
         connect_only: bool | None = None,
+        pel: PEL | None = None,
     ) -> None:
+
+        self.pel: PEL = pel or PEL()
+
         super().__init__(
             *brokers,
             with_real=with_real,
@@ -104,9 +130,7 @@ async def _create_ctx(self) -> AsyncGenerator[list[RedisBroker], None]:
                         wraps=partial(self._fake_connect, broker),
                     ):
                         await broker.connect()
-
                     cluster_stack.enter_context(self._patch_producer(broker))
-
             async with super()._create_ctx() as brokers:
                 yield brokers
 
@@ -116,15 +140,15 @@ def _patch_producer(self, broker: RedisBroker) -> Iterator[None]:
             es.enter_context(
                 change_producer(
                     broker.config.broker_config,
-                    FakeProducer(broker, self.brokers, broker.config),
+                    FakeProducer(broker, self.brokers, broker.config, self.pel),
                 ),
             )
 
             for publisher in cast("list[LogicPublisher]", broker.publishers):
                 es.enter_context(
                     change_producer(
                         publisher,
-                        FakeProducer(broker, self.brokers, publisher.config),
+                        FakeProducer(broker, self.brokers, publisher.config, self.pel),
                     ),
                 )
 
@@ -193,6 +217,7 @@ def __init__(
         broker: RedisBroker,
         brokers: Sequence[RedisBroker],
         config: ParserConfig,
+        pel: PEL | None = None,
     ) -> None:
         self.broker = broker
         self.brokers = brokers
@@ -209,6 +234,7 @@ def __init__(
             default.decode_message,
         )
         self.codec = broker.config.broker_codec or DefaultCodec()
+        self.pel: PEL = pel or PEL()
 
     @property
     def subscribers(self) -> "Iterable[LogicSubscriber]":
@@ -233,20 +259,31 @@ async def publish(self, cmd: "RedisPublishCommand") -> int | bytes:
             serializer=self.broker.config.fd_config._serializer,
             codec=self.codec,
         )
-
         destination = _make_destination_kwargs(cmd)
         visitors = (ChannelVisitor(), ListVisitor(), StreamVisitor())
+        session_id = uuid.uuid4()
 
         for handler in self.subscribers:  # pragma: no branch
             for visitor in visitors:
                 if visited_ch := visitor.visit(**destination, sub=handler):
+                    if pel_entry := self.pel.get_entry(
+                        correlation_id=(cmd.correlation_id, session_id)
+                    ):
+                        await self._execute_handler(
+                            pel_entry.msg, pel_entry.handler, session_id=session_id
+                        )
+                        continue
                     msg = visitor.get_message(
                         visited_ch,
                         body,
                         handler,  # type: ignore[arg-type]
                     )
-
-                    await self._execute_handler(msg, handler)
+                    self.pel.put(
+                        msg=msg,
+                        handler=handler,
+                        correlation_id=(cmd.correlation_id, session_id),
+                    )
+                    await self._execute_handler(msg, handler, session_id=session_id)
 
         return 0
 
@@ -263,18 +300,32 @@ async def request(self, cmd: "RedisPublishCommand") -> "PubSubMessage":
 
         destination = _make_destination_kwargs(cmd)
         visitors = (ChannelVisitor(), ListVisitor(), StreamVisitor())
+        session_id = uuid.uuid4()
 
         for handler in self.subscribers:  # pragma: no branch
             for visitor in visitors:
                 if visited_ch := visitor.visit(**destination, sub=handler):
+                    if pel_entry := self.pel.get_entry(
+                        correlation_id=(cmd.correlation_id, session_id)
+                    ):
+                        await self._execute_handler(
+                            pel_entry.msg, pel_entry.handler, session_id=session_id
+                        )
+                        continue
                     msg = visitor.get_message(
                         visited_ch,
                         body,
                         handler,  # type: ignore[arg-type]
                     )
-
+                    self.pel.put(
+                        msg=msg,
+                        handler=handler,
+                        correlation_id=(cmd.correlation_id, session_id),
+                    )
                     with anyio.fail_after(cmd.timeout):
-                        return await self._execute_handler(msg, handler)
+                        return await self._execute_handler(
+                            msg, handler, session_id=session_id
+                        )
 
         raise SubscriberNotFound
 
@@ -291,30 +342,43 @@ async def publish_batch(self, cmd: "RedisPublishCommand") -> int:
             )
             for m in cmd.batch_bodies
         ]
-
+        session_id = uuid.uuid4()
         visitor = ListVisitor()
         for handler in self.subscribers:  # pragma: no branch
             if visitor.visit(list=cmd.destination, sub=handler):
                 casted_handler = cast("_ListHandlerMixin", handler)
 
                 if casted_handler.list_sub.batch:
+                    if pel_entry := self.pel.get_entry(
+                        correlation_id=(cmd.correlation_id, session_id)
+                    ):
+                        await self._execute_handler(
+                            pel_entry.msg, pel_entry.handler, session_id=session_id
+                        )
+                        continue
                     msg = visitor.get_message(
                         channel=cmd.destination,
                         body=data_to_send,
                         sub=casted_handler,
                     )
-
-                    await self._execute_handler(msg, handler)
+                    self.pel.put(
+                        msg=msg,
+                        handler=handler,
+                        correlation_id=(cmd.correlation_id, session_id),
+                    )
+                    await self._execute_handler(msg, handler, session_id=session_id)
 
         return 0
 
     async def _execute_handler(
         self,
         msg: Any,
         handler: "LogicSubscriber",
+        session_id: uuid.UUID,
     ) -> "PubSubMessage":
         result = await handler.process_message(msg)
-
+        if result.correlation_id:
+            self.pel.remove(correlation_id=(result.correlation_id, session_id))
         return PubSubMessage(
             type="message",
             data=await build_message(

tests/brokers/redis/test_cluster_pubsub_more.py

@@ -44,7 +44,7 @@ async def handler(msg: str) -> None:
                 timeout=self.timeout,
             )
 
-        assert received == ["a", "b"]
+        assert set(received) == {"a", "b"}
 
     @pytest.mark.asyncio()
     async def test_multiple_subscribers_same_channel(