Fixed AttributeError 'drain' for server websocket handler #1613

CHANGES.rst

@@ -1,6 +1,12 @@
 CHANGES
 =======
 
+1.3.1 (2017-02-09)
+------------------
+
+- Fixed AttributeError 'drain' for server websocket handler #1613
+
+
 1.3.0 (2017-02-08)
 ------------------
 

aiohttp/web_reqrep.py

@@ -301,6 +301,11 @@ def transport(self):
         """Transport used for request processing."""
         return self._transport
 
+    @property
+    def transport_pair(self):
+        """Reader and writer used for request processing."""
+        return (self._reader, self._writer)
+
     @reify
     def cookies(self):
         """Return request cookies.

aiohttp/web_ws.py

@@ -105,7 +105,7 @@ def prepare(self, request):
     def _pre_start(self, request):
         try:
             status, headers, parser, writer, protocol = do_handshake(
-                request.method, request.headers, request.transport,
+                request.method, request.headers, request.transport_pair[1],
                 self._protocols)
         except HttpProcessingError as err:
             if err.code == 405:

tests/test_web_websocket_functional.py

@@ -90,6 +90,35 @@ def handler(request):
     assert data['test'] == expected_value
 
 
+@asyncio.coroutine
+def test_websocket_send_drain(loop, test_client):
+    @asyncio.coroutine
+    def handler(request):
+        ws = web.WebSocketResponse()
+        yield from ws.prepare(request)
+
+        ws._writer._limit = 1
+
+        data = yield from ws.receive_json()
+        drain = ws.send_json(data)
+        assert drain
+
+        yield from drain
+        yield from ws.close()
+        return ws
+
+    app = web.Application(loop=loop)
+    app.router.add_route('GET', '/', handler)
+    client = yield from test_client(app)
+
+    ws = yield from client.ws_connect('/')
+    expected_value = 'value'
+    ws.send_json({'test': expected_value})
+
+    data = yield from ws.receive_json()
+    assert data['test'] == expected_value
+
+
 @asyncio.coroutine
 def test_websocket_receive_json(loop, test_client):
     @asyncio.coroutine