diff --git a/tests/test_web_websocket.py b/tests/test_web_websocket.py index cd53eada871..90dfe32ff2c 100644 --- a/tests/test_web_websocket.py +++ b/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() diff --git a/tests/test_web_websocket_functional.py b/tests/test_web_websocket_functional.py index e7c5a5f6d55..67f939aa072 100644 --- a/tests/test_web_websocket_functional.py +++ b/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 diff --git a/tests/test_web_websocket_functional_oldstyle.py b/tests/test_web_websocket_functional_oldstyle.py index cf3dd2bcf8d..3ed40db6188 100644 --- 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 index 3aeb0223cf5..933676619e4 100644 --- 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 index c6992356543..39b20e46509 100644 --- a/tests/test_websocket_client_functional.py +++ b/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):