From 03ca4a3b7873c0d20773b37c195dcd0d23759a51 Mon Sep 17 00:00:00 2001
From: Joongi Kim <me@daybreaker.info>
Date: Wed, 24 Aug 2016 20:17:55 +0900
Subject: [PATCH] Explicitly return TimeoutError on WebSocketResponse.close()
 (#1084) (#1089)

---
 aiohttp/web_ws.py                      |  4 ++++
 tests/test_web_websocket_functional.py | 25 +++++++++++++++----------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py
index bbecbdcf479..e458c5690cb 100644
--- a/aiohttp/web_ws.py
+++ b/aiohttp/web_ws.py
@@ -205,6 +205,10 @@ def close(self, *, code=1000, message=b''):
                 if msg.type == WSMsgType.CLOSE:
                     self._close_code = msg.data
                     return True
+
+            self._close_code = 1006
+            self._exception = asyncio.TimeoutError()
+            return True
         else:
             return False
 
diff --git a/tests/test_web_websocket_functional.py b/tests/test_web_websocket_functional.py
index 28011f1b4eb..7f872fb5432 100644
--- a/tests/test_web_websocket_functional.py
+++ b/tests/test_web_websocket_functional.py
@@ -4,6 +4,7 @@
 
 import aiohttp
 from aiohttp import helpers, web
+from aiohttp._ws_impl import WSMsgType
 
 
 @asyncio.coroutine
@@ -218,7 +219,7 @@ def handler(request):
 
 @asyncio.coroutine
 def test_close_timeout(create_app_and_client, loop):
-    closed = helpers.create_future(loop)
+    aborted = helpers.create_future(loop)
 
     @asyncio.coroutine
     def handler(request):
@@ -227,12 +228,14 @@ def handler(request):
         assert 'request' == (yield from ws.receive_str())
         ws.send_str('reply')
         begin = ws._loop.time()
-        yield from ws.close()
+        assert (yield from ws.close())
         elapsed = ws._loop.time() - begin
         assert elapsed < 0.201, \
             'close() should have returned before ' \
             'at most 2x timeout.'
-        closed.set_result(1)
+        assert ws.close_code == 1006
+        assert isinstance(ws.exception(), asyncio.TimeoutError)
+        aborted.set_result(1)
         return ws
 
     app, client = yield from create_app_and_client()
@@ -245,15 +248,17 @@ def handler(request):
     # The server closes here.  Then the client sends bogus messages with an
     # internval shorter than server-side close timeout, to make the server
     # hanging indefinitely.
-    yield from asyncio.sleep(0.04, loop=loop)
+    yield from asyncio.sleep(0.08, loop=loop)
+    msg = yield from ws._reader.read()
+    assert msg.type == WSMsgType.CLOSE
     ws.send_str('hang')
-    yield from asyncio.sleep(0.04, loop=loop)
+    yield from asyncio.sleep(0.08, loop=loop)
     ws.send_str('hang')
-    yield from asyncio.sleep(0.04, loop=loop)
+    yield from asyncio.sleep(0.08, loop=loop)
     ws.send_str('hang')
-    yield from asyncio.sleep(0.04, loop=loop)
+    yield from asyncio.sleep(0.08, loop=loop)
     ws.send_str('hang')
-    yield from asyncio.sleep(0.04, loop=loop)
-    # The server should have been closed now.
-    assert 1 == (yield from closed)
+    yield from asyncio.sleep(0.08, loop=loop)
+    assert (yield from aborted)
+
     yield from ws.close()