Added stream broker, updated redis.

Author: s3rius

pyproject.toml

@@ -55,6 +55,7 @@ warn_return_any = false
 [[tool.mypy.overrides]]
 module = ['redis']
 ignore_missing_imports = true
+ignore_errors = true
 strict = false
 
 [build-system]

taskiq_redis/redis_backend.py

@@ -34,8 +34,8 @@
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _Redis: TypeAlias = Redis[bytes]
-    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+    _Redis: TypeAlias = Redis[bytes]  # type: ignore
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
 else:
     _Redis: TypeAlias = Redis
     _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
@@ -258,7 +258,7 @@ def __init__(
         :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
             and result_px_time are equal zero.
         """
-        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+        self.redis: "RedisCluster" = RedisCluster.from_url(
             redis_url,
             **connection_kwargs,
         )
@@ -275,14 +275,10 @@ def __init__(
             ),
         )
         if unavailable_conditions:
-            raise ExpireTimeMustBeMoreThanZeroError(
-                "You must select one expire time param and it must be more than zero.",
-            )
+            raise ExpireTimeMustBeMoreThanZeroError
 
         if self.result_ex_time and self.result_px_time:
-            raise DuplicateExpireTimeSelectedError(
-                "Choose either result_ex_time or result_px_time.",
-            )
+            raise DuplicateExpireTimeSelectedError
 
     def _task_name(self, task_id: str) -> str:
         if self.prefix_str is None:
@@ -291,7 +287,7 @@ def _task_name(self, task_id: str) -> str:
 
     async def shutdown(self) -> None:
         """Closes redis connection."""
-        await self.redis.aclose()  # type: ignore[attr-defined]
+        await self.redis.aclose()
         await super().shutdown()
 
     async def set_result(
@@ -327,7 +323,7 @@ async def is_result_ready(self, task_id: str) -> bool:
 
         :returns: True if the result is ready else False.
         """
-        return bool(await self.redis.exists(self._task_name(task_id)))  # type: ignore[attr-defined]
+        return bool(await self.redis.exists(self._task_name(task_id)))
 
     async def get_result(
         self,
@@ -344,11 +340,11 @@ async def get_result(
         """
         task_name = self._task_name(task_id)
         if self.keep_results:
-            result_value = await self.redis.get(  # type: ignore[attr-defined]
+            result_value = await self.redis.get(
                 name=task_name,
             )
         else:
-            result_value = await self.redis.getdel(  # type: ignore[attr-defined]
+            result_value = await self.redis.getdel(
                 name=task_name,
             )
 
@@ -400,7 +396,7 @@ async def get_progress(
         :param task_id: task's id.
         :return: task's TaskProgress instance.
         """
-        result_value = await self.redis.get(  # type: ignore[attr-defined]
+        result_value = await self.redis.get(
             name=self._task_name(task_id) + PROGRESS_KEY_SUFFIX,
         )
 
@@ -465,14 +461,10 @@ def __init__(
             ),
         )
         if unavailable_conditions:
-            raise ExpireTimeMustBeMoreThanZeroError(
-                "You must select one expire time param and it must be more than zero.",
-            )
+            raise ExpireTimeMustBeMoreThanZeroError
 
         if self.result_ex_time and self.result_px_time:
-            raise DuplicateExpireTimeSelectedError(
-                "Choose either result_ex_time or result_px_time.",
-            )
+            raise DuplicateExpireTimeSelectedError
 
     def _task_name(self, task_id: str) -> str:
         if self.prefix_str is None:
@@ -610,4 +602,4 @@ async def get_progress(
     async def shutdown(self) -> None:
         """Shutdown sentinel connections."""
         for sentinel in self.sentinel.sentinels:
-            await sentinel.aclose()  # type: ignore[attr-defined]
+            await sentinel.aclose()

taskiq_redis/redis_broker.py

@@ -28,7 +28,7 @@
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
 else:
     _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
 
@@ -122,7 +122,7 @@ async def kick(self, message: BrokerMessage) -> None:
         """
         queue_name = message.labels.get("queue_name") or self.queue_name
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.lpush(queue_name, message.message)
+            await redis_conn.lpush(queue_name, message.message)  # type: ignore
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -137,7 +137,7 @@ async def listen(self) -> AsyncGenerator[bytes, None]:
         while True:
             try:
                 async with Redis(connection_pool=self.connection_pool) as redis_conn:
-                    yield (await redis_conn.brpop(self.queue_name))[
+                    yield (await redis_conn.brpop(self.queue_name))[  # type: ignore
                         redis_brpop_data_position
                     ]
             except ConnectionError as exc:
@@ -238,7 +238,7 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
                     self.consumer_name,
                     {
                         self.queue_name: ">",
-                        **self.additional_streams,
+                        **self.additional_streams,  # type: ignore
                     },
                     block=self.block,
                     noack=False,

taskiq_redis/redis_cluster_broker.py

@@ -30,7 +30,7 @@ def __init__(
         """
         super().__init__()
 
-        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+        self.redis: "RedisCluster[bytes]" = RedisCluster.from_url(  # type: ignore
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
@@ -40,7 +40,7 @@ def __init__(
 
     async def shutdown(self) -> None:
         """Closes redis connection pool."""
-        await self.redis.aclose()  # type: ignore[attr-defined]
+        await self.redis.aclose()
         await super().shutdown()
 
 
@@ -55,7 +55,7 @@ async def kick(self, message: BrokerMessage) -> None:
 
         :param message: message to append.
         """
-        await self.redis.lpush(self.queue_name, message.message)  # type: ignore[attr-defined]
+        await self.redis.lpush(self.queue_name, message.message)  # type: ignore
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -68,7 +68,7 @@ async def listen(self) -> AsyncGenerator[bytes, None]:
         """
         redis_brpop_data_position = 1
         while True:
-            value = await self.redis.brpop([self.queue_name])  # type: ignore[attr-defined]
+            value = await self.redis.brpop([self.queue_name])  # type: ignore
             yield value[redis_brpop_data_position]
 
 
@@ -155,7 +155,7 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
                 self.consumer_name,
                 {
                     self.queue_name: ">",
-                    **self.additional_streams,
+                    **self.additional_streams,  # type: ignore
                 },
                 block=self.block,
                 noack=False,

taskiq_redis/redis_sentinel_broker.py

@@ -27,7 +27,7 @@
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _Redis: TypeAlias = Redis[bytes]
+    _Redis: TypeAlias = Redis[bytes]  # type: ignore
 else:
     _Redis: TypeAlias = Redis
 
@@ -117,7 +117,7 @@ async def kick(self, message: BrokerMessage) -> None:
         """
         queue_name = message.labels.get("queue_name") or self.queue_name
         async with self._acquire_master_conn() as redis_conn:
-            await redis_conn.lpush(queue_name, message.message)
+            await redis_conn.lpush(queue_name, message.message)  # type: ignore
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -131,7 +131,7 @@ async def listen(self) -> AsyncGenerator[bytes, None]:
         redis_brpop_data_position = 1
         async with self._acquire_master_conn() as redis_conn:
             while True:
-                yield (await redis_conn.brpop(self.queue_name))[
+                yield (await redis_conn.brpop(self.queue_name))[  # type: ignore
                     redis_brpop_data_position
                 ]
 
@@ -226,7 +226,7 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
                     self.consumer_name,
                     {
                         self.queue_name: ">",
-                        **self.additional_streams,
+                        **self.additional_streams,  # type: ignore
                     },
                     block=self.block,
                     noack=False,

taskiq_redis/schedule_source.py

@@ -21,8 +21,8 @@
     from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    _Redis: TypeAlias = Redis[bytes]
-    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+    _Redis: TypeAlias = Redis[bytes]  # type: ignore
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
 else:
     _Redis: TypeAlias = Redis
     _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
@@ -140,7 +140,7 @@ def __init__(
         **connection_kwargs: Any,
     ) -> None:
         self.prefix = prefix
-        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+        self.redis: "RedisCluster" = RedisCluster.from_url(
             url,
             **connection_kwargs,
         )
@@ -150,7 +150,7 @@ def __init__(
 
     async def delete_schedule(self, schedule_id: str) -> None:
         """Remove schedule by id."""
-        await self.redis.delete(f"{self.prefix}:{schedule_id}")  # type: ignore[attr-defined]
+        await self.redis.delete(f"{self.prefix}:{schedule_id}")
 
     async def add_schedule(self, schedule: ScheduledTask) -> None:
         """
@@ -159,7 +159,7 @@ async def add_schedule(self, schedule: ScheduledTask) -> None:
         :param schedule: schedule to add.
         :param schedule_id: schedule id.
         """
-        await self.redis.set(  # type: ignore[attr-defined]
+        await self.redis.set(
             f"{self.prefix}:{schedule.schedule_id}",
             self.serializer.dumpb(model_dump(schedule)),
         )
@@ -173,8 +173,8 @@ async def get_schedules(self) -> List[ScheduledTask]:
         :return: list of schedules.
         """
         schedules = []
-        async for key in self.redis.scan_iter(f"{self.prefix}:*"):  # type: ignore[attr-defined]
-            raw_schedule = await self.redis.get(key)  # type: ignore[attr-defined]
+        async for key in self.redis.scan_iter(f"{self.prefix}:*"):
+            raw_schedule = await self.redis.get(key)
             parsed_schedule = model_validate(
                 ScheduledTask,
                 self.serializer.loadb(raw_schedule),
@@ -189,7 +189,7 @@ async def post_send(self, task: ScheduledTask) -> None:
 
     async def shutdown(self) -> None:
         """Shut down the schedule source."""
-        await self.redis.aclose()  # type: ignore[attr-defined]
+        await self.redis.aclose()
 
 
 class RedisSentinelScheduleSource(ScheduleSource):
@@ -288,4 +288,4 @@ async def post_send(self, task: ScheduledTask) -> None:
     async def shutdown(self) -> None:
         """Shut down the schedule source."""
         for sentinel in self.sentinel.sentinels:
-            await sentinel.aclose()  # type: ignore[attr-defined]
+            await sentinel.aclose()

tests/test_broker.py

@@ -14,6 +14,7 @@
     RedisStreamClusterBroker,
     RedisStreamSentinelBroker,
 )
+from taskiq_redis.redis_broker import RedisStreamBroker
 
 
 def test_no_url_should_raise_typeerror() -> None:
@@ -129,6 +130,44 @@ async def test_list_queue_broker(
     await broker.shutdown()
 
 
+@pytest.mark.anyio
+async def test_stream_broker(
+    valid_broker_message: BrokerMessage,
+    redis_url: str,
+) -> None:
+    """
+    Test that messages are published and read correctly by ListQueueBroker.
+
+    We create two workers that listen and send a message to them.
+    Expect only one worker to receive the same message we sent.
+    """
+    broker = RedisStreamBroker(
+        url=redis_url,
+        queue_name=uuid.uuid4().hex,
+        consumer_group_name=uuid.uuid4().hex,
+    )
+    await broker.startup()
+
+    worker1_task = asyncio.create_task(get_message(broker))
+    worker2_task = asyncio.create_task(get_message(broker))
+
+    await broker.kick(valid_broker_message)
+
+    await asyncio.wait(
+        [worker1_task, worker2_task],
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+
+    assert worker1_task.done() != worker2_task.done()
+    message = worker1_task.result() if worker1_task.done() else worker2_task.result()
+    assert isinstance(message, AckableMessage)
+    assert message.data == valid_broker_message.message
+    await message.ack()  # type: ignore
+    worker1_task.cancel()
+    worker2_task.cancel()
+    await broker.shutdown()
+
+
 @pytest.mark.anyio
 async def test_list_queue_broker_max_connections(
     valid_broker_message: BrokerMessage,
@@ -196,14 +235,13 @@ async def test_stream_cluster_broker(
         consumer_group_name=uuid.uuid4().hex,
     )
     await broker.startup()
+
     worker_task = asyncio.create_task(get_message(broker))
-    await asyncio.sleep(0.3)
 
     await broker.kick(valid_broker_message)
-    await asyncio.sleep(0.3)
 
-    assert worker_task.done()
-    result = worker_task.result()
+    result = await worker_task
+
     assert isinstance(result, AckableMessage)
     assert result.data == valid_broker_message.message
     await result.ack()  # type: ignore
@@ -291,13 +329,10 @@ async def test_streams_sentinel_broker(
     )
     await broker.startup()
     worker_task = asyncio.create_task(get_message(broker))
-    await asyncio.sleep(0.3)
 
     await broker.kick(valid_broker_message)
-    await asyncio.sleep(0.3)
 
-    assert worker_task.done()
-    result = worker_task.result()
+    result = await worker_task
     assert isinstance(result, AckableMessage)
     assert result.data == valid_broker_message.message
     await result.ack()  # type: ignore