def test_send_bytes_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_bytes(b'bytes')
def test_send_json_closed(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
yield from ws.close()
with pytest.raises(RuntimeError):
ws.send_json({'type': 'json'})
def test_pong_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.pong()
def test_close_idempotent(make_request, writer):
req = make_request("GET", "/")
ws = WebSocketResponse()
yield from ws.prepare(req)
assert (yield from ws.close(code=1, message="message1"))
assert ws.closed
assert not (yield from ws.close(code=2, message="message2"))
def test_pong_closed(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
yield from ws.close()
with pytest.raises(RuntimeError):
ws.pong()
def test_send_str_closed(make_request):
req = make_request("GET", "/")
ws = WebSocketResponse()
yield from ws.prepare(req)
yield from ws.close()
with pytest.raises(RuntimeError):
ws.send_str("string")
def test_send_json_closed(make_request):
req = make_request("GET", "/")
ws = WebSocketResponse()
yield from ws.prepare(req)
yield from ws.close()
with pytest.raises(RuntimeError):
ws.send_json({"type": "json"})
def test_close_idempotent(make_request, writer):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
assert (yield from ws.close(code=1, message='message1'))
assert ws.closed
assert not (yield from ws.close(code=2, message='message2'))
def wsHandler(request):
resp = WebSocketResponse()
ok, protocol = resp.can_start(request)
if not ok:
with open(WS_FILE, 'rb') as fp:
return Response(body=fp.read(), content_type='text/html')
yield from resp.prepare(request)
print('Someone joined.')
for ws in request.app['sockets']:
ws.send_str('Someone joined')
request.app['sockets'].append(resp)
while True:
msg = yield from resp.receive()
if msg.tp == MsgType.text:
obj = json.loads(msg.data)
retVal = parseRequest(obj)
for ws in request.app['sockets']:
if ws is resp:
ws.send_str(json.dumps(retVal))
else:
break
request.app['sockets'].remove(resp)
print('Someone disconnected.')
for ws in request.app['sockets']:
ws.send_str('Someone disconnected.')
return resp
def test_concurrent_receive(self):
req = self.make_request("GET", "/")
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
ws._waiting = True
self.assertRaises(RuntimeError, self.loop.run_until_complete, ws.receive())
def test_prepare_twice_idempotent(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
impl1 = yield from ws.prepare(req)
impl2 = yield from ws.prepare(req)
assert impl1 is impl2
def wshandler(request):
resp = WebSocketResponse()
ok, protocol = resp.can_start(request)
if not ok:
with open(WS_FILE, "rb") as fp:
return Response(body=fp.read(), content_type="text/html")
yield from resp.prepare(request)
print("Someone joined.")
for ws in request.app["sockets"]:
ws.send_str("Someone joined")
request.app["sockets"].append(resp)
while True:
msg = yield from resp.receive()
if msg.tp == MsgType.text:
for ws in request.app["sockets"]:
if ws is not resp:
ws.send_str(msg.data)
else:
break
request.app["sockets"].remove(resp)
print("Someone disconnected.")
for ws in request.app["sockets"]:
ws.send_str("Someone disconnected.")
return resp
def test_close_idempotent(make_request, writer):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
ws._reader.feed_data(WS_CLOSED_MESSAGE, 0)
assert (yield from ws.close(code=1, message='message1'))
assert ws.closed
assert not (yield from ws.close(code=2, message='message2'))
def test_concurrent_receive(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
ws._waiting = True
with pytest.raises(RuntimeError):
yield from ws.receive()
def test_can_prepare_started(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
with pytest.raises(RuntimeError) as ctx:
ws.can_prepare(req)
assert 'Already started' in str(ctx.value)
def test_write_eof_idempotent(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
yield from ws.close()
yield from ws.write_eof()
yield from ws.write_eof()
yield from ws.write_eof()
def test_close_idempotent(self):
req = self.make_request("GET", "/")
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
writer = mock.Mock()
ws._writer = writer
self.assertTrue(self.loop.run_until_complete(ws.close(code=1, message="message1")))
self.assertTrue(ws.closed)
self.assertFalse(self.loop.run_until_complete(ws.close(code=2, message="message2")))
def test_multiple_receive_on_close_connection(self):
req = self.make_request("GET", "/")
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
self.loop.run_until_complete(ws.close())
self.loop.run_until_complete(ws.receive())
self.loop.run_until_complete(ws.receive())
self.loop.run_until_complete(ws.receive())
self.loop.run_until_complete(ws.receive())
self.assertRaises(RuntimeError, self.loop.run_until_complete, ws.receive())
def test_receive_timeouterror(self):
req = self.make_request("GET", "/")
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
res = asyncio.Future(loop=self.loop)
res.set_exception(asyncio.TimeoutError())
ws._reader.read.return_value = res
self.assertRaises(asyncio.TimeoutError, self.loop.run_until_complete, ws.receive())
def test_pong_closed(make_request, mocker):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
ws._reader.feed_data(WS_CLOSED_MESSAGE, 0)
yield from ws.close()
mocker.spy(ws_logger, 'warning')
ws.pong()
assert ws_logger.warning.called
def test_write_eof_idempotent(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
ws._reader.feed_data(WS_CLOSED_MESSAGE, 0)
yield from ws.close()
yield from ws.write_eof()
yield from ws.write_eof()
yield from ws.write_eof()
def test_receive_cancelled(make_request, loop, reader):
req = make_request("GET", "/")
ws = WebSocketResponse()
yield from ws.prepare(req)
res = helpers.create_future(loop)
res.set_exception(asyncio.CancelledError())
reader.read = make_mocked_coro(res)
with pytest.raises(asyncio.CancelledError):
yield from ws.receive()
def test_receive_timeouterror(make_request, loop, reader):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
res = helpers.create_future(loop)
res.set_exception(asyncio.TimeoutError())
reader.read = make_mocked_coro(res)
with pytest.raises(asyncio.TimeoutError):
yield from ws.receive()
def test_receive_cancelled(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
res = asyncio.Future(loop=self.loop)
res.set_exception(asyncio.CancelledError())
ws._reader.read.return_value = res
self.assertRaises(
asyncio.CancelledError,
self.loop.run_until_complete, ws.receive())
def go():
req = self.make_request("GET", "/")
ws = WebSocketResponse()
yield from ws.prepare(req)
@asyncio.coroutine
def receive():
return websocket.Message(websocket.MSG_BINARY, b"data", b"")
ws.receive = receive
with self.assertRaises(TypeError):
yield from ws.receive_str()
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(TypeError):
yield from ws.receive_bytes()
def test_receive_str_nonstring(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
@asyncio.coroutine
def receive():
return WSMessage(WSMsgType.BINARY, b'data', b'')
ws.receive = receive
with pytest.raises(TypeError):
yield from ws.receive_str()
def test_write_eof_idempotent(self):
req = self.make_request("GET", "/")
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
self.loop.run_until_complete(ws.close())
@asyncio.coroutine
def go():
yield from ws.write_eof()
yield from ws.write_eof()
yield from ws.write_eof()
self.loop.run_until_complete(go())
def test_receive_bytes_nonsbytes(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
@asyncio.coroutine
def receive():
return WSMessage(WSMsgType.TEXT, 'data', b'')
ws.receive = receive
with pytest.raises(TypeError):
yield from ws.receive_bytes()
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(TypeError):
yield from ws.receive_bytes()
def test_receive_bytes_nonsbytes(make_request):
req = 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 pytest.raises(TypeError):
yield from ws.receive_bytes()
def test_multiple_receive_on_close_connection(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
yield from ws.close()
yield from ws.receive()
yield from ws.receive()
yield from ws.receive()
yield from ws.receive()
with pytest.raises(RuntimeError):
yield from ws.receive()
def test_multiple_receive_on_close_connection(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
ws._reader.feed_data(CLOSED_MESSAGE, 0)
yield from ws.close()
yield from ws.receive()
yield from ws.receive()
yield from ws.receive()
yield from ws.receive()
with pytest.raises(RuntimeError):
yield from ws.receive()
def test_receive_exc_in_reader(make_request, loop, reader):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
exc = ValueError()
res = helpers.create_future(loop)
res.set_exception(exc)
reader.read = make_mocked_coro(res)
msg = yield from ws.receive()
assert msg.tp == WSMsgType.error
assert msg.data is exc
assert ws.exception() is exc
def function113(function271):
var2027 = function271('GET', '/')
var2897 = WebSocketResponse()
yield from var2897.prepare(var2027)
var962 = ValueError()
var2897._writer = mock.Mock()
var2897._writer.close.side_effect = var962
yield from var2897.close()
assert var2897.closed
assert (var2897.exception() is var962)
var2897._closed = False
var2897._writer.close.side_effect = asyncio.CancelledError()
with pytest.raises(asyncio.CancelledError):
yield from var2897.close()
def wrapper(request):
params = request.GET
kwargs = {}
if authenticate:
kwargs["user"] = handle_auth(params.get("token", None))
ws = WebSocketResponse()
try:
yield from ws.prepare(request)
yield from func(ws, params, **kwargs)
except Exception as err: # pragma: no cover
logger.error(str(err))
return ws
def test_receive_client_disconnected(make_request, loop, reader):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
exc = errors.ClientDisconnectedError()
res = helpers.create_future(loop)
res.set_exception(exc)
reader.read = make_mocked_coro(res)
msg = yield from ws.receive()
assert ws.closed
assert msg.tp == WSMsgType.close
assert msg.data is None
assert ws.exception() is None
def test_close_exc2(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
exc = ValueError()
self.writer.close.side_effect = exc
ws._writer = self.writer
self.loop.run_until_complete(ws.close())
self.assertTrue(ws.closed)
self.assertIs(ws.exception(), exc)
ws._closed = False
self.writer.close.side_effect = asyncio.CancelledError()
self.assertRaises(asyncio.CancelledError, self.loop.run_until_complete,
ws.close())
def notifications(request, response):
notifications = NotificationManager.instance()
ws = WebSocketResponse()
yield from ws.prepare(request)
asyncio.async(process_websocket(ws))
with notifications.queue() as queue:
while True:
try:
notification = yield from queue.get_json(5)
except asyncio.futures.CancelledError:
break
if ws.closed:
break
ws.send_str(notification)
return ws
def function2536(function271, arg68):
var3402 = function271('GET', '/')
var3771 = WebSocketResponse()
yield from var3771.prepare(var3402)
var3771._reader = mock.Mock()
var2011 = ValueError()
var1794 = helpers.create_future(arg68)
var1794.set_exception(var2011)
var3771._reader.read = make_mocked_coro(var1794)
var3771._payload_writer.drain = mock.Mock()
var3771._payload_writer.drain.return_value = helpers.create_future(arg68)
var3771._payload_writer.drain.return_value.set_result(True)
var2571 = yield from var3771.function1900()
assert (var2571.type == WSMsgType.ERROR)
assert (var2571.type is var2571.tp)
assert (var2571.data is var2011)
assert (var3771.exception() is var2011)
def test_close_exc2(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
exc = ValueError()
ws._writer = mock.Mock()
ws._writer.close.side_effect = exc
yield from ws.close()
assert ws.closed
assert ws.exception() is exc
ws._closed = False
ws._writer.close.side_effect = asyncio.CancelledError()
with pytest.raises(asyncio.CancelledError):
yield from ws.close()
def test_receive_exc_in_reader(make_request, loop, reader):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
exc = ValueError()
res = helpers.create_future(loop)
res.set_exception(exc)
reader.read = make_mocked_coro(res)
ws._resp_impl.transport.drain.return_value = helpers.create_future(loop)
ws._resp_impl.transport.drain.return_value.set_result(True)
msg = yield from ws.receive()
assert msg.type == WSMsgType.ERROR
assert msg.type is msg.tp
assert msg.data is exc
assert ws.exception() is exc
def test_receive_exc_in_reader(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
exc = ValueError()
res = helpers.create_future(self.loop)
res.set_exception(exc)
ws._reader.read.return_value = res
@asyncio.coroutine
def go():
msg = yield from ws.receive()
self.assertTrue(msg.tp, MsgType.error)
self.assertIs(msg.data, exc)
self.assertIs(ws.exception(), exc)
self.loop.run_until_complete(go())
def test_receive_client_disconnected(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
exc = errors.ClientDisconnectedError()
res = helpers.create_future(self.loop)
res.set_exception(exc)
ws._reader.read.return_value = res
@asyncio.coroutine
def go():
msg = yield from ws.receive()
self.assertTrue(ws.closed)
self.assertTrue(msg.tp, MsgType.close)
self.assertIs(msg.data, None)
self.assertIs(ws.exception(), None)
self.loop.run_until_complete(go())
def socket(request):
resp = WebSocketResponse()
yield from resp.prepare(request)
if 'amqp' not in request.app:
transport, protocol = yield from aioamqp.connect()
request.app['amqp'] = protocol
@asyncio.coroutine
def cleanup(app):
yield from protocol.close(timeout=1.0)
transport.close()
request.app.register_on_finish(cleanup)
amqp = request.app['amqp']
channel = yield from amqp.channel()
yield from channel.exchange_declare(exchange_name='demo-room',
type_name='fanout')
result = yield from channel.queue_declare('', exclusive=True)
yield from channel.queue_bind(result['queue'], 'demo-room', routing_key='')
# Consume messages from the queue
@asyncio.coroutine
def message(channel, body, envelope, properties):
if not resp.closed:
resp.send_str(body.decode('utf-8'))
yield from channel.basic_client_ack(envelope.delivery_tag)
yield from channel.basic_consume(message, queue_name=result['queue'])
# Broadcast messages to the queue
while True:
msg = yield from resp.receive()
if msg.tp == MsgType.text:
yield from channel.publish(msg.data,
exchange_name='demo-room',
routing_key='')
else:
break
# Client requested close
yield from channel.close()
return resp
def test_close_exc(make_request, reader, loop):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
exc = ValueError()
reader.read.return_value = helpers.create_future(loop)
reader.read.return_value.set_exception(exc)
yield from ws.close()
assert ws.closed
assert ws.exception() is exc
ws._closed = False
reader.read.return_value = helpers.create_future(loop)
reader.read.return_value.set_exception(asyncio.CancelledError())
with pytest.raises(asyncio.CancelledError):
yield from ws.close()
assert ws.close_code == 1006
def function479(function271, arg926, arg363):
var1796 = function271('GET', '/')
var300 = WebSocketResponse()
yield from var300.prepare(var1796)
var300._reader = mock.Mock()
var4357 = ValueError()
var300._reader.read.return_value = helpers.create_future(arg926)
var300._reader.read.return_value.set_exception(var4357)
var300._payload_writer.drain = mock.Mock()
var300._payload_writer.drain.return_value = helpers.create_future(arg926)
var300._payload_writer.drain.return_value.set_result(True)
yield from var300.close()
assert var300.closed
assert (var300.exception() is var4357)
var300._closed = False
var300._reader.read.return_value = helpers.create_future(arg926)
var300._reader.read.return_value.set_exception(asyncio.CancelledError())
with pytest.raises(asyncio.CancelledError):
yield from var300.close()
assert (var300.close_code == 1006)
def test_close_exc(self):
req = self.make_request('GET', '/')
reader = self.reader.set_parser.return_value = mock.Mock()
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
exc = ValueError()
reader.read.return_value = helpers.create_future(self.loop)
reader.read.return_value.set_exception(exc)
self.loop.run_until_complete(ws.close())
self.assertTrue(ws.closed)
self.assertIs(ws.exception(), exc)
ws._closed = False
reader.read.return_value = helpers.create_future(self.loop)
reader.read.return_value.set_exception(asyncio.CancelledError())
self.assertRaises(asyncio.CancelledError, self.loop.run_until_complete,
ws.close())
self.assertEqual(ws.close_code, 1006)
def wrapper(request):
params = request.GET
kwargs = {}
token = params.get("token", None)
if authenticate:
kwargs["user"] = handle_auth(params.get("token", None))
redis_ = yield from get_async_redis()
manager = SubscriptionManager(redis_)
kwargs["manager"] = manager
ws = WebSocketResponse()
try:
yield from ws.prepare(request)
yield from func(ws, params, **kwargs)
except Exception as err: # pragma: no cover
logger.error(str(err))
finally:
yield from manager.stop()
return ws
def test_close_exc(make_request, loop, mocker):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
ws._reader = mock.Mock()
exc = ValueError()
ws._reader.read.return_value = helpers.create_future(loop)
ws._reader.read.return_value.set_exception(exc)
ws._payload_writer.drain = mock.Mock()
ws._payload_writer.drain.return_value = helpers.create_future(loop)
ws._payload_writer.drain.return_value.set_result(True)
yield from ws.close()
assert ws.closed
assert ws.exception() is exc
ws._closed = False
ws._reader.read.return_value = helpers.create_future(loop)
ws._reader.read.return_value.set_exception(asyncio.CancelledError())
with pytest.raises(asyncio.CancelledError):
yield from ws.close()
assert ws.close_code == 1006
def function1287(function271):
var1020 = function271('GET', '/', headers=CIMultiDict({}))
var1629 = WebSocketResponse()
with pytest.raises(HTTPBadRequest):
yield from var1629.prepare(var1020)
def test_prepare_invalid_method(make_request):
req = make_request('POST', '/')
ws = WebSocketResponse()
with pytest.raises(HTTPMethodNotAllowed):
yield from ws.prepare(req)
def test_prepare_without_upgrade(make_request):
req = make_request('GET', '/', headers=CIMultiDict({}))
ws = WebSocketResponse()
with pytest.raises(HTTPBadRequest):
yield from ws.prepare(req)
def test_send_bytes_nonbytes(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
with pytest.raises(TypeError):
ws.send_bytes('string')
def test_send_json_nonjson(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
with pytest.raises(TypeError):
ws.send_json(set())
def test_start_without_upgrade(self):
req = self.make_request('GET', '/', headers=CIMultiDict({}))
ws = WebSocketResponse()
with self.assertRaises(HTTPBadRequest):
self.loop.run_until_complete(ws.prepare(req))
def test_start_invalid_method(self):
req = self.make_request('POST', '/')
ws = WebSocketResponse()
with self.assertRaises(HTTPMethodNotAllowed):
self.loop.run_until_complete(ws.prepare(req))
def test_can_prepare_started(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
with self.assertRaisesRegex(RuntimeError, 'Already started'):
ws.can_prepare(req)
def test_send_bytes_nonbytes(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
self.loop.run_until_complete(ws.prepare(req))
with self.assertRaises(TypeError):
ws.send_bytes('string')
def function265(function271):
var2045 = function271('GET', '/')
var4391 = WebSocketResponse()
var4138 = yield from var4391.prepare(var2045)
var1947 = yield from var4391.prepare(var2045)
assert (var4138 is var1947)
