Fix for (#976): add test for websocket send_json and receive_json

tests/test_web_websocket.py

@@ -1,5 +1,6 @@
 import asyncio
 import unittest
+import json
 from unittest import mock
 from aiohttp import CIMultiDict, helpers
 from aiohttp.web import (
@@ -65,6 +66,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 +96,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
@@ -128,6 +144,25 @@ def receive():
 
         self.loop.run_until_complete(go())
 
+    def test_receive_json_nonjson(self):
+
+        @asyncio.coroutine
+        def go():
+            req = self.make_request('GET', '/')
+            ws = WebSocketResponse()
+            yield from ws.prepare(req)
+
+            @asyncio.coroutine
+            def receive():
+                return websocket.Message(websocket.MSG_TEXT, 'data', b'')
+
+            ws.receive = receive
+
+            with self.assertRaises(json.decoder.JSONDecodeError):
+                yield from ws.receive_json()
+
+        self.loop.run_until_complete(go())
+
     def test_send_str_nonstring(self):
         req = self.make_request('GET', '/')
         ws = WebSocketResponse()
@@ -142,6 +177,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 +238,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()

tests/test_web_websocket_functional.py

@@ -61,6 +61,30 @@ def handler(request):
     assert payload in resp.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
 def test_websocket_receive_json(create_app_and_client):
     @asyncio.coroutine

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)

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')

tests/test_websocket_client_functional.py

@@ -52,6 +52,30 @@ def handler(request):
     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()
+
+
 @pytest.mark.run_loop
 def test_ping_pong(create_app_and_client, loop):