Make connection handling more explicit and remove outdated examples

redis_websocket_api/handler.py

@@ -132,7 +132,7 @@ def channel_is_allowed(self, channel_name):
     async def create(cls, redis, websocket, read_timeout=None):
         """Create a handler instance setting up tasks and queues."""
         self = cls(redis, websocket, read_timeout=read_timeout)
-        self.consumer_task = asyncio.ensure_future(self._websocket_reader())
+        self.consumer_task = asyncio.create_task(self._websocket_reader())
         return self
 
     async def listen(self):
@@ -141,7 +141,7 @@ async def listen(self):
         This coroutine blocks for up to self.read_timeout seconds.
         """
         await self._send("websocket", {"status": "open"})
-        while not self.consumer_task.done():
+        while not self.consumer_task.done() and not self.websocket.closed:
             try:
                 await self._queue_reader()
             except asyncio.TimeoutError:
@@ -150,9 +150,10 @@ async def listen(self):
 
     async def close(self):
         """Close all connections and cancel all tasks."""
-        asyncio.wait_for(  # attempt to say goodbye to the client
-            self.websocket.close(), self.read_timeout
-        )
+        if self.websocket.open:
+            await asyncio.wait_for(  # attempt to say goodbye to the client
+                self.websocket.close(), self.read_timeout
+            )
         self.consumer_task.cancel()
 
 

redis_websocket_api/protocol.py

@@ -16,7 +16,7 @@ class CommandsMixin:
 
     The commend name is translated to a method name like this:
 
-        "_handle_{name}_command".format(name=a_tralis_protocol_command.lower())
+        "_handle_{name}_command".format(name=a_custom_protocol_command.lower())
     """
 
     async def _handle_del_command(self, channel_name):

redis_websocket_api/server.py

@@ -49,13 +49,25 @@ async def websocket_handler(self, websocket, path):
 
         logger.info("Client %s connected", websocket.remote_address)
         handler = await self.handler_class.create(
-            self.redis, websocket, read_timeout=self.read_timeout,
+            self.redis,
+            websocket,
+            read_timeout=self.read_timeout,
         )
         self.handlers[websocket.remote_address] = handler
+        handler_listen_task = asyncio.create_task(
+            asyncio.wait_for(handler.listen(), self.keep_alive_timeout)
+        )
         try:
-            await asyncio.wait_for(handler.listen(), self.keep_alive_timeout)
+            await asyncio.wait(
+                {
+                    handler_listen_task,
+                    websocket.handler_task,
+                },
+                return_when="FIRST_COMPLETED",
+            )
         finally:
             del self.handlers[websocket.remote_address]
+            handler_listen_task.cancel()
             await handler.close()
             logger.info("Client %s removed", websocket.remote_address)
 

tests/conftest.py

@@ -1,3 +1,4 @@
+import asyncio
 from unittest.mock import MagicMock
 
 import pytest
@@ -37,6 +38,7 @@ def __getattr__(self, name):
 def websocket():
     websocket = AsyncMagicMock()
     websocket.remote_address = ("EGG", 2000)
+    websocket.handler_task = asyncio.sleep(0)
     return websocket
 
 

tests/test_server.py

@@ -9,8 +9,7 @@ def test_websocket_handler_creation(loop, server, websocket):
     server.handlers = MagicMock()
     websocket.await_recv.side_effect = exceptions.ConnectionClosed(1001, "foo")
 
-    with pytest.warns(RuntimeWarning):
-        asyncio.run(server.websocket_handler(websocket, "/foo"))
+    asyncio.run(server.websocket_handler(websocket, "/foo"))
 
     assert websocket.await_recv.call_count == 1
     assert websocket.await_send.call_count == 1