Websockets: Stop iteration when connection closes. (#1145)

* Websockets: Stop iteration when connection closes.

See issue #1144.

* WebsocketResponse: Async for now stops iteration when WSMsgType.Closed is passed down (and not just WsMsgType.Close).

* Tests: Added tests for the async for functionality of the websocket response and websocket client response.

* Moved python3.5 'async for closed' tests unto the test_py35 directory equivalents.

* Added Frederik Peter Aalund to CONTRIBUTORS.txt

* Updated CHANGES.rst to reflect that #1144 is now fixed.

CHANGES.rst

@@ -57,7 +57,7 @@ CHANGES
 
 -
 
--
+- Websockets: Stop `async for` iteration when connection is closed #1144
 
 - Ensure TestClient HTTP methods return a context manager #1318
 

CONTRIBUTORS.txt

@@ -53,6 +53,7 @@ Erich Healy
 Eugene Chernyshov
 Eugene Naydenov
 Frederik Gladhorn
+Frederik Peter Aalund
 Gabriel Tremblay
 Gennady Andreyev
 Georges Dubus

aiohttp/client_ws.py

@@ -188,6 +188,6 @@ def __aiter__(self):
         @asyncio.coroutine
         def __anext__(self):
             msg = yield from self.receive()
-            if msg.type == WSMsgType.CLOSE:
+            if msg.type == WSMsgType.CLOSE or msg.type == WSMsgType.CLOSED:
                 raise StopAsyncIteration  # NOQA
             return msg

aiohttp/web_ws.py

@@ -315,6 +315,6 @@ def __aiter__(self):
         @asyncio.coroutine
         def __anext__(self):
             msg = yield from self.receive()
-            if msg.type == WSMsgType.CLOSE:
+            if msg.type == WSMsgType.CLOSE or msg.type == WSMsgType.CLOSED:
                 raise StopAsyncIteration  # NOQA
             return msg

tests/test_py35/test_client_websocket_35.py

@@ -1,7 +1,7 @@
 import pytest
 
 import aiohttp
-from aiohttp import web
+from aiohttp import helpers, web
 
 async def test_client_ws_async_for(loop, test_client):
     items = ['q1', 'q2', 'q3']
@@ -74,3 +74,38 @@ async def handler(request):
             assert msg.data == 'request/answer'
 
         assert ws.closed
+
+
+async def test_closed_async_for(loop, test_client):
+
+    closed = helpers.create_future(loop)
+
+    async def handler(request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+
+        try:
+            ws.send_bytes(b'started')
+            await ws.receive_bytes()
+        finally:
+            closed.set_result(1)
+        return ws
+
+    app = web.Application(loop=loop)
+    app.router.add_route('GET', '/', handler)
+    client = await test_client(app)
+    resp = await client.ws_connect('/')
+
+    messages = []
+    async for msg in resp:
+        messages.append(msg)
+        if b'started' == msg.data:
+            resp.send_bytes(b'ask')
+            await resp.close()
+
+    assert 1 == len(messages)
+    assert messages[0].type == aiohttp.WSMsgType.BINARY
+    assert messages[0].data == b'started'
+    assert resp.closed
+
+    await closed

tests/test_py35/test_web_websocket_35.py

@@ -1,5 +1,6 @@
 import aiohttp
 from aiohttp import helpers, web
+from aiohttp._ws_impl import WSMsgType
 
 async def test_server_ws_async_for(loop, test_server):
     closed = helpers.create_future(loop)
@@ -29,3 +30,39 @@ async def handler(request):
 
     await resp.close()
     await closed
+
+
+async def test_closed_async_for(loop, test_client):
+
+    closed = helpers.create_future(loop)
+
+    async def handler(request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+
+        messages = []
+        async for msg in ws:
+            messages.append(msg)
+            if 'stop' == msg.data:
+                ws.send_str('stopping')
+                await ws.close()
+
+        assert 1 == len(messages)
+        assert messages[0].type == WSMsgType.TEXT
+        assert messages[0].data == 'stop'
+
+        closed.set_result(None)
+        return ws
+
+    app = web.Application(loop=loop)
+    app.router.add_get('/', handler)
+    client = await test_client(app)
+
+    ws = await client.ws_connect('/')
+    ws.send_str('stop')
+    msg = await ws.receive()
+    assert msg.type == WSMsgType.TEXT
+    assert msg.data == 'stopping'
+
+    await ws.close()
+    await closed