Async Connection: Allow PubSub.run() without previous subscribe()

redis/asyncio/client.py

@@ -693,16 +693,24 @@ async def execute_command(self, *args: EncodableT):
         # legitimate message off the stack if the connection is already
         # subscribed to one or more channels
 
+        await self.connect()
+        connection = self.connection
+        kwargs = {"check_health": not self.subscribed}
+        await self._execute(connection, connection.send_command, *args, **kwargs)
+
+    async def connect(self):
+        """
+        Ensure that the PubSub is connected
+        """
         if self.connection is None:
             self.connection = await self.connection_pool.get_connection(
                 "pubsub", self.shard_hint
             )
             # register a callback that re-subscribes to any channels we
             # were listening to when we were disconnected
             self.connection.register_connect_callback(self.on_connect)
-        connection = self.connection
-        kwargs = {"check_health": not self.subscribed}
-        await self._execute(connection, connection.send_command, *args, **kwargs)
+        else:
+            await self.connection.connect()
 
     async def _disconnect_raise_connect(self, conn, error):
         """
@@ -962,6 +970,7 @@ async def run(
             if handler is None:
                 raise PubSubError(f"Pattern: '{pattern}' has no handler registered")
 
+        await self.connect()
         while True:
             try:
                 await self.get_message(

tests/test_asyncio/test_pubsub.py

@@ -2,6 +2,7 @@
 import sys
 from typing import Optional
 
+import async_timeout
 import pytest
 
 if sys.version_info[0:2] == (3, 6):
@@ -658,3 +659,35 @@ def exception_handler_callback(e, pubsub) -> None:
             except asyncio.CancelledError:
                 pass
         assert str(e) == "error"
+
+    async def test_late_subscribe(self, r: redis.Redis):
+        def callback(message):
+            messages.put_nowait(message)
+
+        messages = asyncio.Queue()
+        p = r.pubsub()
+        task = asyncio.get_event_loop().create_task(p.run())
+        # wait until loop gets settled.  Add a subscription
+        await asyncio.sleep(0.1)
+        await p.subscribe(foo=callback)
+        # wait tof the subscribe to finish.  Cannot use _subscribe() because
+        # p.run() is already accepting messages
+        await asyncio.sleep(0.1)
+        await r.publish("foo", "bar")
+        message = None
+        try:
+            async with async_timeout.timeout(0.1):
+                message = await messages.get()
+        except asyncio.TimeoutError:
+            pass
+        task.cancel()
+        # we expect a cancelled error, not the Runtime error
+        # ("did you forget to call subscribe()"")
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert message == {
+            "channel": b"foo",
+            "data": b"bar",
+            "pattern": None,
+            "type": "message",
+        }