Work on tests

Author: asvetlov

CHANGES.txt

@@ -18,6 +18,9 @@ CHANGES
 - Avoid using of mutable CIMultiDict kw param in make_mocked_request
   #997
 
+- Make WebSocketResponse.close a little bit faster by avoiding new
+  task creating just for timeout measurement
+
 0.22.2 (07-23-2016)
 -------------------
 

aiohttp/web_ws.py

@@ -3,7 +3,7 @@
 import json
 import warnings
 
-from . import hdrs
+from . import hdrs, Timeout
 from .errors import HttpProcessingError, ClientDisconnectedError
 from .websocket import do_handshake, Message, WebSocketError
 from .websocket_client import MsgType, closedMessage
@@ -186,9 +186,9 @@ def close(self, *, code=1000, message=b''):
 
             while True:
                 try:
-                    msg = yield from asyncio.wait_for(
-                        self._reader.read(),
-                        timeout=self._timeout, loop=self._loop)
+                    with Timeout(timeout=self._timeout,
+                                 loop=self._loop):
+                        msg = yield from self._reader.read()
                 except asyncio.CancelledError:
                     self._close_code = 1006
                     raise

tests/test_web_websocket_functional.py

@@ -2,7 +2,9 @@
 
 import asyncio
 import pytest
+import aiohttp
 from aiohttp import web
+from aiohttp import helpers
 
 
 @pytest.mark.run_loop
@@ -85,3 +87,38 @@ def handler(request):
 
     resp = yield from ws.receive()
     assert resp.data == expected_value
+
+
+@pytest.mark.run_loop
+def test_send_recv_text(create_app_and_client, loop):
+
+    closed = helpers.create_future(loop)
+
+    @asyncio.coroutine
+    def handler(request):
+        ws = web.WebSocketResponse()
+        yield from ws.prepare(request)
+        msg = yield from ws.receive_str()
+        ws.send_str(msg+'/answer')
+        yield from ws.close()
+        closed.set_result(1)
+        return ws
+
+    app, client = yield from create_app_and_client()
+    app.router.add_route('GET', '/', handler)
+
+    ws = yield from client.ws_connect('/')
+    ws.send_str('ask')
+    msg = yield from ws.receive()
+    assert msg.tp == aiohttp.MsgType.text
+    assert 'ask/answer' == msg.data
+
+    msg = yield from ws.receive()
+    assert msg.tp == aiohttp.MsgType.close
+    assert msg.data == 1000
+    assert msg.extra == ''
+
+    assert ws.closed
+    assert ws.close_code == 1000
+
+    yield from closed

tests/test_web_websocket_functional_oldstyle.py

@@ -82,41 +82,6 @@ def connect_ws(self, url, protocol=None):
 
         return response, reader, writer
 
-    def test_send_recv_text(self):
-
-        closed = helpers.create_future(self.loop)
-
-        @asyncio.coroutine
-        def handler(request):
-            ws = web.WebSocketResponse()
-            yield from ws.prepare(request)
-            msg = yield from ws.receive_str()
-            ws.send_str(msg+'/answer')
-            yield from ws.close()
-            closed.set_result(1)
-            return ws
-
-        @asyncio.coroutine
-        def go():
-            _, _, url = yield from self.create_server('GET', '/', handler)
-            resp, reader, writer = yield from self.connect_ws(url)
-            writer.send('ask')
-            msg = yield from reader.read()
-            self.assertEqual(msg.tp, websocket.MSG_TEXT)
-            self.assertEqual('ask/answer', msg.data)
-
-            msg = yield from reader.read()
-            self.assertEqual(msg.tp, websocket.MSG_CLOSE)
-            self.assertEqual(msg.data, 1000)
-            self.assertEqual(msg.extra, '')
-
-            writer.close()
-
-            yield from closed
-            resp.close()
-
-        self.loop.run_until_complete(go())
-
     def test_send_recv_bytes(self):
 
         closed = helpers.create_future(self.loop)