Add unittest for PubSub.connect() (#2167)

* Add unittest for PubSub reconnect * fix linting

Author: kristjanvalur

tests/test_asyncio/test_pubsub.py

@@ -1,4 +1,5 @@
 import asyncio
+import functools
 import sys
 from typing import Optional
 
@@ -20,6 +21,18 @@
 pytestmark = pytest.mark.asyncio(forbid_global_loop=True)
 
 
+def with_timeout(t):
+ def wrapper(corofunc):
+ @functools.wraps(corofunc)
+ async def run(*args, **kwargs):
+ async with async_timeout.timeout(t):
+ return await corofunc(*args, **kwargs)
+
+ return run
+
+ return wrapper
+
+
 async def wait_for_message(pubsub, timeout=0.1, ignore_subscribe_messages=False):
  now = asyncio.get_event_loop().time()
  timeout = now + timeout
@@ -603,6 +616,75 @@ async def test_get_message_with_timeout_returns_none(self, r: redis.Redis):
  assert await p.get_message(timeout=0.01) is None
 
 
+@pytest.mark.onlynoncluster
+class TestPubSubReconnect:
+ # @pytest.mark.xfail
+ @with_timeout(2)
+ async def test_reconnect_listen(self, r: redis.Redis):
+ """
+ Test that a loop processing PubSub messages can survive
+ a disconnect, by issuing a connect() call.
+ """
+ messages = asyncio.Queue()
+ pubsub = r.pubsub()
+ interrupt = False
+
+ async def loop():
+ # must make sure the task exits
+ async with async_timeout.timeout(2):
+ nonlocal interrupt
+ await pubsub.subscribe("foo")
+ while True:
+ # print("loop")
+ try:
+ try:
+ await pubsub.connect()
+ await loop_step()
+ # print("succ")
+ except redis.ConnectionError:
+ await asyncio.sleep(0.1)
+ except asyncio.CancelledError:
+ # we use a cancel to interrupt the "listen"
+ # when we perform a disconnect
+ # print("cancel", interrupt)
+ if interrupt:
+ interrupt = False
+ else:
+ raise
+
+ async def loop_step():
+ # get a single message via listen()
+ async for message in pubsub.listen():
+ await messages.put(message)
+ break
+
+ task = asyncio.get_event_loop().create_task(loop())
+ # get the initial connect message
+ async with async_timeout.timeout(1):
+ message = await messages.get()
+ assert message == {
+ "channel": b"foo",
+ "data": 1,
+ "pattern": None,
+ "type": "subscribe",
+ }
+ # now, disconnect the connection.
+ await pubsub.connection.disconnect()
+ interrupt = True
+ task.cancel() # interrupt the listen call
+ # await another auto-connect message
+ message = await messages.get()
+ assert message == {
+ "channel": b"foo",
+ "data": 1,
+ "pattern": None,
+ "type": "subscribe",
+ }
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+
+
 @pytest.mark.onlynoncluster
 class TestPubSubRun:
  async def _subscribe(self, p, *args, **kwargs):