Repository: aio-libs/aiohttp

Issue:
Fix websocket json support
Implement server `resp.send_json()` and client `resp.send_json()`/`resp.receive_json()`



Discussion:

Resolving patch:
diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py
--- a/aiohttp/web_ws.py
+++ b/aiohttp/web_ws.py
@@ -154,6 +154,9 @@ def send_bytes(self, data):
                             type(data))
         self._writer.send(data, binary=True)
 
+    def send_json(self, data, *, dumps=json.dumps):
+        self.send_str(dumps(data))
+
     @asyncio.coroutine
     def wait_closed(self):  # pragma: no cover
         warnings.warn(
@@ -288,12 +291,8 @@ def receive_bytes(self):
 
     @asyncio.coroutine
     def receive_json(self, *, loads=json.loads):
-        msg = yield from self.receive()
-        if msg.tp != MsgType.text:
-            raise TypeError(
-                "Received message {}:{!r} is not str".format(msg.tp, msg.data)
-            )
-        return msg.json(loads=loads)
+        data = yield from self.receive_str()
+        return loads(data)
 
     def write(self, data):
         raise RuntimeError("Cannot call .write() for websocket")
diff --git a/aiohttp/websocket_client.py b/aiohttp/websocket_client.py
--- a/aiohttp/websocket_client.py
+++ b/aiohttp/websocket_client.py
@@ -3,6 +3,7 @@
 import asyncio
 
 import sys
+import json
 from enum import IntEnum
 
 from .websocket import Message
@@ -88,6 +89,9 @@ def send_bytes(self, data):
                             type(data))
         self._writer.send(data, binary=True)
 
+    def send_json(self, data, *, dumps=json.dumps):
+        self.send_str(dumps(data))
+
     @asyncio.coroutine
     def close(self, *, code=1000, message=b''):
         if not self._closed:
@@ -171,6 +175,28 @@ def receive(self):
         finally:
             self._waiting = False
 
+    @asyncio.coroutine
+    def receive_str(self):
+        msg = yield from self.receive()
+        if msg.tp != MsgType.text:
+            raise TypeError(
+                "Received message {}:{!r} is not str".format(msg.tp, msg.data))
+        return msg.data
+
+    @asyncio.coroutine
+    def receive_bytes(self):
+        msg = yield from self.receive()
+        if msg.tp != MsgType.binary:
+            raise TypeError(
+                "Received message {}:{!r} is not bytes".format(msg.tp,
+                                                               msg.data))
+        return msg.data
+
+    @asyncio.coroutine
+    def receive_json(self, *, loads=json.loads):
+        data = yield from self.receive_str()
+        return loads(data)
+
     if PY_35:
         @asyncio.coroutine
         def __aiter__(self):


Test patch:
diff --git a/tests/test_web_websocket.py b/tests/test_web_websocket.py
--- a/tests/test_web_websocket.py
+++ b/tests/test_web_websocket.py
@@ -65,6 +65,11 @@ def test_nonstarted_send_bytes(self):
         with self.assertRaises(RuntimeError):
             ws.send_bytes(b'bytes')
 
+    def test_nonstarted_send_json(self):
+        ws = WebSocketResponse()
+        with self.assertRaises(RuntimeError):
+            ws.send_json({'type': 'json'})
+
     def test_nonstarted_close(self):
         ws = WebSocketResponse()
         with self.assertRaises(RuntimeError):
@@ -90,6 +95,16 @@ def go():
 
         self.loop.run_until_complete(go())
 
+    def test_nonstarted_receive_json(self):
+
+        @asyncio.coroutine
+        def go():
+            ws = WebSocketResponse()
+            with self.assertRaises(RuntimeError):
+                yield from ws.receive_json()
+
+        self.loop.run_until_complete(go())
+
     def test_receive_str_nonstring(self):
 
         @asyncio.coroutine
@@ -142,6 +157,13 @@ def test_send_bytes_nonbytes(self):
         with self.assertRaises(TypeError):
             ws.send_bytes('string')
 
+    def test_send_json_nonjson(self):
+        req = self.make_request('GET', '/')
+        ws = WebSocketResponse()
+        self.loop.run_until_complete(ws.prepare(req))
+        with self.assertRaises(TypeError):
+            ws.send_json(set())
+
     def test_write(self):
         ws = WebSocketResponse()
         with self.assertRaises(RuntimeError):
@@ -196,6 +218,14 @@ def test_send_bytes_closed(self):
         with self.assertRaises(RuntimeError):
             ws.send_bytes(b'bytes')
 
+    def test_send_json_closed(self):
+        req = self.make_request('GET', '/')
+        ws = WebSocketResponse()
+        self.loop.run_until_complete(ws.prepare(req))
+        self.loop.run_until_complete(ws.close())
+        with self.assertRaises(RuntimeError):
+            ws.send_json({'type': 'json'})
+
     def test_ping_closed(self):
         req = self.make_request('GET', '/')
         ws = WebSocketResponse()
diff --git a/tests/test_web_websocket_functional.py b/tests/test_web_websocket_functional.py
--- a/tests/test_web_websocket_functional.py
+++ b/tests/test_web_websocket_functional.py
@@ -38,14 +38,12 @@ def test_websocket_json_invalid_message(create_app_and_client):
     def handler(request):
         ws = web.WebSocketResponse()
         yield from ws.prepare(request)
-        msg = yield from ws.receive()
-
         try:
-            msg.json()
+            yield from ws.receive_json()
         except ValueError:
-            ws.send_str("ValueError raised: '%s'" % msg.data)
+            ws.send_str('ValueError was raised')
         else:
-            raise Exception("No ValueError was raised")
+            raise Exception('No Exception')
         finally:
             yield from ws.close()
         return ws
@@ -57,8 +55,32 @@ def handler(request):
     payload = 'NOT A VALID JSON STRING'
     ws.send_str(payload)
 
-    resp = yield from ws.receive()
-    assert payload in resp.data
+    data = yield from ws.receive_str()
+    assert 'ValueError was raised' in data
+
+
+@pytest.mark.run_loop
+def test_websocket_send_json(create_app_and_client):
+    @asyncio.coroutine
+    def handler(request):
+        ws = web.WebSocketResponse()
+        yield from ws.prepare(request)
+
+        data = yield from ws.receive_json()
+        ws.send_json(data)
+
+        yield from ws.close()
+        return ws
+
+    app, client = yield from create_app_and_client()
+    app.router.add_route('GET', '/', handler)
+
+    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
 
 
 @pytest.mark.run_loop
diff --git a/tests/test_web_websocket_functional_oldstyle.py b/tests/test_web_websocket_functional_oldstyle.py
--- a/tests/test_web_websocket_functional_oldstyle.py
+++ b/tests/test_web_websocket_functional_oldstyle.py
@@ -152,6 +152,41 @@ def go():
 
         self.loop.run_until_complete(go())
 
+    def test_send_recv_json(self):
+        closed = helpers.create_future(self.loop)
+
+        @asyncio.coroutine
+        def handler(request):
+            ws = web.WebSocketResponse()
+            yield from ws.prepare(request)
+            data = yield from ws.receive_json()
+            ws.send_json({'response': data['request']})
+            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('{"request": "test"}')
+            msg = yield from reader.read()
+            data = msg.json()
+            self.assertEqual(msg.tp, websocket.MSG_TEXT)
+            self.assertEqual(data['response'], 'test')
+
+            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_auto_pong_with_closing_by_peer(self):
 
         closed = helpers.create_future(self.loop)
diff --git a/tests/test_websocket_client.py b/tests/test_websocket_client.py
--- a/tests/test_websocket_client.py
+++ b/tests/test_websocket_client.py
@@ -334,6 +334,7 @@ def test_send_data_after_close(self, m_req, m_os, WebSocketWriter):
         self.assertRaises(RuntimeError, resp.pong)
         self.assertRaises(RuntimeError, resp.send_str, 's')
         self.assertRaises(RuntimeError, resp.send_bytes, b'b')
+        self.assertRaises(RuntimeError, resp.send_json, {})
 
     @mock.patch('aiohttp.client.WebSocketWriter')
     @mock.patch('aiohttp.client.os')
@@ -357,6 +358,7 @@ def test_send_data_type_errors(self, m_req, m_os, WebSocketWriter):
 
         self.assertRaises(TypeError, resp.send_str, b's')
         self.assertRaises(TypeError, resp.send_bytes, 'b')
+        self.assertRaises(TypeError, resp.send_json, set())
 
     @mock.patch('aiohttp.client.WebSocketWriter')
     @mock.patch('aiohttp.client.os')
diff --git a/tests/test_websocket_client_functional.py b/tests/test_websocket_client_functional.py
--- a/tests/test_websocket_client_functional.py
+++ b/tests/test_websocket_client_functional.py
@@ -22,8 +22,8 @@ def handler(request):
     resp = yield from client.ws_connect('/')
     resp.send_str('ask')
 
-    msg = yield from resp.receive()
-    assert msg.data == 'ask/answer'
+    data = yield from resp.receive_str()
+    assert data == 'ask/answer'
     yield from resp.close()
 
 
@@ -46,9 +46,33 @@ def handler(request):
 
     resp.send_bytes(b'ask')
 
-    msg = yield from resp.receive()
-    assert msg.data == b'ask/answer'
+    data = yield from resp.receive_bytes()
+    assert data == b'ask/answer'
+
+    yield from resp.close()
+
+
+@pytest.mark.run_loop
+def test_send_recv_json(create_app_and_client):
+
+    @asyncio.coroutine
+    def handler(request):
+        ws = web.WebSocketResponse()
+        yield from ws.prepare(request)
+
+        data = yield from ws.receive_json()
+        ws.send_json({'response': data['request']})
+        yield from ws.close()
+        return ws
+
+    app, client = yield from create_app_and_client()
+    app.router.add_route('GET', '/', handler)
+    resp = yield from client.ws_connect('/')
+    payload = {'request': 'test'}
+    resp.send_json(payload)
 
+    data = yield from resp.receive_json()
+    assert data['response'] == payload['request']
     yield from resp.close()