From 0f74eae4123d8d7baa4b5982474545b031aa8251 Mon Sep 17 00:00:00 2001
From: Frederik Aalund <frederikaalund@gmail.com>
Date: Fri, 21 Oct 2016 18:15:44 +0200
Subject: [PATCH] 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                                 |  2 +-
 CONTRIBUTORS.txt                            |  1 +
 aiohttp/client_ws.py                        |  2 +-
 aiohttp/web_ws.py                           |  2 +-
 tests/test_py35/test_client_websocket_35.py | 37 ++++++++++++++++++++-
 tests/test_py35/test_web_websocket_35.py    | 37 +++++++++++++++++++++
 6 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index e45adf0661b..0c3907da8be 100644
--- a/CHANGES.rst
+++ b/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
 
diff --git a/CONTRIBUTORS.txt b/CONTRIBUTORS.txt
index db3fdd66a0e..05846ce0f75 100644
--- a/CONTRIBUTORS.txt
+++ b/CONTRIBUTORS.txt
@@ -53,6 +53,7 @@ Erich Healy
 Eugene Chernyshov
 Eugene Naydenov
 Frederik Gladhorn
+Frederik Peter Aalund
 Gabriel Tremblay
 Gennady Andreyev
 Georges Dubus
diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py
index 984669c10a0..6b09e0300eb 100644
--- a/aiohttp/client_ws.py
+++ b/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
diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py
index 8873225b2a3..596a85aa7d9 100644
--- a/aiohttp/web_ws.py
+++ b/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
diff --git a/tests/test_py35/test_client_websocket_35.py b/tests/test_py35/test_client_websocket_35.py
index 7eaa8b7d212..55e9f605079 100644
--- a/tests/test_py35/test_client_websocket_35.py
+++ b/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
diff --git a/tests/test_py35/test_web_websocket_35.py b/tests/test_py35/test_web_websocket_35.py
index c804cf9cb41..db34e31f2fa 100644
--- a/tests/test_py35/test_web_websocket_35.py
+++ b/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