async def get(self):
ws = WebSocketResponse()
await ws.prepare(self.request)
with self.resource:
while not ws.closed:
try:
item = await asyncio.wait_for(self.resource.queue.get(), 2)
except asyncio.TimeoutError:
ws.ping()
continue
except asyncio.CancelledError:
break
if item['action'] == 'cleared':
data = {'action': 'cleared'}
elif item['action'] == 'received':
data = {
'action': 'received',
'msg': serialize_message(item['msg']),
}
else:
data = None
if data:
ws.send_str(json.dumps(data))
return ws
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)
assert (yield from ws.close(code=1, message="message1"))
assert ws.closed
assert not (yield from ws.close(code=2, message="message2"))
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_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 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_start_twice_idempotent(self):
req = self.make_request("GET", "/")
ws = WebSocketResponse()
with self.assertWarns(DeprecationWarning):
impl1 = ws.start(req)
impl2 = ws.start(req)
self.assertIs(impl1, impl2)
async def ws_handler(request):
channel = request.match_info['channel']
ws = WebSocketResponse()
await ws.prepare(request)
print('Someone joined channel : %s.' % channel)
try:
request.app['sockets'][channel].append(ws)
except KeyError:
request.app['sockets'][channel] = []
request.app['sockets'][channel].append(ws)
while not ws.closed:
msg = await ws.receive()
if msg.tp == MsgType.text:
if msg.data == 'close':
await ws.close()
else:
print('The client responded: ' + msg.data)
elif msg.tp == MsgType.close:
print('websocket connection closed')
elif msg.tp == MsgType.error:
print('ws connection closed with exception %s' % ws.exception())
request.app['sockets'][channel].remove(ws)
print('Someone disconnected from channel : %s.' % channel)
return ws
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_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_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_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_start_twice_idempotent(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
with pytest.warns(DeprecationWarning):
impl1 = ws.start(req)
impl2 = ws.start(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:
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_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 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')
resp.start(request)
print('Someone joined.')
for ws in request.app['sockets']:
ws.send_str('Someone joined')
request.app['sockets'].append(resp)
try:
while True:
msg = yield from resp.receive_str()
print(msg)
for ws in request.app['sockets']:
if ws is not resp:
ws.send_str(msg)
except WSClientDisconnectedError:
if resp not in request.app['sockets']:
return resp
request.app['sockets'].remove(resp)
print('Someone disconnected.')
for ws in request.app['sockets']:
ws.send_str('Someone disconnected.')
raise
async def wshandler(request):
resp = WebSocketResponse()
ok, protocol = resp.can_prepare(request)
if not ok:
with open(WS_FILE, 'rb') as fp:
return Response(body=fp.read(), content_type='text/html')
await resp.prepare(request)
try:
print('Someone joined.')
for ws in request.app['sockets']:
await ws.send_str('Someone joined')
request.app['sockets'].append(resp)
async for msg in resp:
if msg.type == WSMsgType.TEXT:
for ws in request.app['sockets']:
if ws is not resp:
await ws.send_str(msg.data)
else:
return resp
return resp
finally:
request.app['sockets'].remove(resp)
print('Someone disconnected.')
for ws in request.app['sockets']:
await ws.send_str('Someone disconnected.')
def wshandler(request):
print("wshandler", request)
resp = WebSocketResponse()
ok, protocol = resp.can_start(request)
if not ok:
print("http")
with open(WS_FILE, 'rb') as fp:
return Response(body=fp.read(), content_type='text/html')
print("ws")
resp.start(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()
async def test_can_prepare_started(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
await ws.prepare(req)
with pytest.raises(RuntimeError) as ctx:
ws.can_prepare(req)
assert 'Already started' in str(ctx.value)
def test_concurrent_receive(self):
req = self.make_request('GET', '/')
ws = WebSocketResponse()
ws.start(req)
ws._waiting = True
self.assertRaises(
RuntimeError, self.loop.run_until_complete, ws.receive())
async def test_concurrent_receive(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
await ws.prepare(req)
ws._waiting = True
with pytest.raises(RuntimeError):
await 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_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_can_prepare_invalid_method(make_request):
req = make_request('POST', '/')
ws = WebSocketResponse()
assert (False, None) == ws.can_prepare(req)
def test_send_str_nonstring(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
yield from ws.prepare(req)
with pytest.raises(TypeError):
ws.send_str(b'bytes')
def test_nonstarted_receive_bytes():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
yield from ws.receive_bytes()
async def test_nonstarted_close():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.close()
async def test_nonstarted_send_bytes():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.send_bytes(b'bytes')
async def test_wait_closed_before_start():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.close()
async def test_prepare_invalid_method(make_request):
req = make_request('POST', '/')
ws = WebSocketResponse()
with pytest.raises(HTTPMethodNotAllowed):
await ws.prepare(req)
def test_nonstarted_send_str():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
ws.send_str('string')
def test_nonstarted_pong():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
ws.pong()
async def lobby_socket_handler(request):
users = request.app["users"]
sockets = request.app["lobbysockets"]
seeks = request.app["seeks"]
ws = WebSocketResponse(heartbeat=3.0, receive_timeout=10.0)
ws_ready = ws.can_prepare(request)
if not ws_ready.ok:
raise web.HTTPFound("/")
await ws.prepare(request)
session = await aiohttp_session.get_session(request)
session_user = session.get("user_name")
user = users[session_user] if session_user is not None and session_user in users else None
log.debug("-------------------------- NEW lobby WEBSOCKET by %s" % user)
try:
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
if msg.data == "close":
log.debug("Got 'close' msg.")
break
else:
data = json.loads(msg.data)
if not data["type"] == "pong":
log.debug("Websocket (%s) message: %s" % (id(ws), msg))
if data["type"] == "get_seeks":
response = get_seeks(seeks)
await ws.send_json(response)
elif data["type"] == "create_ai_challenge":
no = await is_playing(request, user, ws)
if no:
continue
variant = data["variant"]
engine = users.get("Fairy-Stockfish")
if engine is None or not engine.online():
# TODO: message that engine is offline, but capture BOT will play instead
engine = users.get("Random-Mover")
seek = Seek(
user, variant,
fen=data["fen"],
color=data["color"],
base=data["minutes"],
inc=data["increment"],
byoyomi_period=data["byoyomi_period"],
level=data["level"],
rated=data["rated"],
chess960=data["chess960"],
handicap=data["handicap"])
# print("SEEK", user, variant, data["fen"], data["color"], data["minutes"], data["increment"], data["level"], False, data["chess960"])
seeks[seek.id] = seek
response = await new_game(request.app, engine, seek.id)
await ws.send_json(response)
if response["type"] != "error":
gameId = response["gameId"]
engine.game_queues[gameId] = asyncio.Queue()
await engine.event_queue.put(challenge(seek, response))
elif data["type"] == "create_seek":
no = await is_playing(request, user, ws)
if no:
continue
print("create_seek", data)
create_seek(seeks, user, data, ws)
await lobby_broadcast(sockets, get_seeks(seeks))
if data.get("target"):
queue = users[data["target"]].notify_queue
if queue is not None:
await queue.put(json.dumps({"notify": "new_challenge"}))
elif data["type"] == "delete_seek":
del seeks[data["seekID"]]
del user.seeks[data["seekID"]]
await lobby_broadcast(sockets, get_seeks(seeks))
elif data["type"] == "accept_seek":
no = await is_playing(request, user, ws)
if no:
continue
if data["seekID"] not in seeks:
continue
seek = seeks[data["seekID"]]
# print("accept_seek", seek.as_json)
response = await new_game(request.app, user, data["seekID"])
await ws.send_json(response)
if seek.user.bot:
gameId = response["gameId"]
seek.user.game_queues[gameId] = asyncio.Queue()
await seek.user.event_queue.put(challenge(seek, response))
else:
await seek.ws.send_json(response)
# Inform others, new_game() deleted accepted seek allready.
await lobby_broadcast(sockets, get_seeks(seeks))
elif data["type"] == "lobby_user_connected":
if session_user is not None:
if data["username"] and data["username"] != session_user:
log.info("+++ Existing lobby_user %s socket connected as %s." % (session_user, data["username"]))
session_user = data["username"]
if session_user in users:
user = users[session_user]
else:
user = User(request.app, username=data["username"], anon=data["username"].startswith("Anon-"))
users[user.username] = user
response = {"type": "lobbychat", "user": "", "message": "%s joined the lobby" % session_user}
else:
if session_user in users:
user = users[session_user]
else:
user = User(request.app, username=data["username"], anon=data["username"].startswith("Anon-"))
users[user.username] = user
response = {"type": "lobbychat", "user": "", "message": "%s joined the lobby" % session_user}
else:
log.info("+++ Existing lobby_user %s socket reconnected." % data["username"])
session_user = data["username"]
if session_user in users:
user = users[session_user]
else:
user = User(request.app, username=data["username"], anon=data["username"].startswith("Anon-"))
users[user.username] = user
response = {"type": "lobbychat", "user": "", "message": "%s rejoined the lobby" % session_user}
await lobby_broadcast(sockets, response)
# update websocket
user.lobby_sockets.add(ws)
sockets[user.username] = user.lobby_sockets
response = {"type": "lobby_user_connected", "username": user.username}
await ws.send_json(response)
response = {"type": "fullchat", "lines": list(request.app["chat"])}
await ws.send_json(response)
# send game count
response = {"type": "g_cnt", "cnt": request.app["g_cnt"]}
await ws.send_json(response)
# send user count
if len(user.game_sockets) == 0:
# not connected to any game socket but connected to lobby socket
if len(user.lobby_sockets) == 1:
request.app["u_cnt"] += 1
response = {"type": "u_cnt", "cnt": request.app["u_cnt"]}
await lobby_broadcast(sockets, response)
else:
response = {"type": "u_cnt", "cnt": request.app["u_cnt"]}
await ws.send_json(response)
elif data["type"] == "lobbychat":
response = {"type": "lobbychat", "user": user.username, "message": data["message"]}
await lobby_broadcast(sockets, response)
request.app["chat"].append(response)
elif data["type"] == "logout":
await ws.close()
elif data["type"] == "disconnect":
# Used only to test socket disconnection...
await ws.close(code=1009)
elif msg.type == aiohttp.WSMsgType.CLOSED:
log.debug("--- Lobby websocket %s msg.type == aiohttp.WSMsgType.CLOSED" % id(ws))
break
elif msg.type == aiohttp.WSMsgType.ERROR:
log.error("--- Lobby ws %s msg.type == aiohttp.WSMsgType.ERROR" % id(ws))
break
else:
log.debug("--- Lobby ws other msg.type %s %s" % (msg.type, msg))
except Exception as e:
log.error("!!! Lobby ws exception occured: %s" % type(e))
finally:
log.debug("---fianlly: await ws.close()")
await ws.close()
if user is not None:
if ws in user.lobby_sockets:
user.lobby_sockets.remove(ws)
# online user counter will be updated in quit_lobby also!
if len(user.lobby_sockets) == 0:
if user.username in sockets:
del sockets[user.username]
# not connected to lobby socket and not connected to game socket
if len(user.game_sockets) == 0:
request.app["u_cnt"] -= 1
response = {"type": "u_cnt", "cnt": request.app["u_cnt"]}
await lobby_broadcast(sockets, response)
response = {"type": "lobbychat", "user": "", "message": "%s left the lobby" % user.username}
await lobby_broadcast(sockets, response)
await user.clear_seeks(sockets, seeks)
return ws
def test_nonstarted_close():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
yield from ws.close()
def test_can_start_ok(make_request):
req = make_request('GET', '/', protocols=True)
ws = WebSocketResponse(protocols=('chat',))
with pytest.warns(DeprecationWarning):
assert (True, 'chat') == ws.can_start(req)
async def test_nonstarted_pong() -> None:
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.pong()
def test_can_prepare_without_upgrade(make_request):
req = make_request('GET', '/', headers=CIMultiDict({}))
ws = WebSocketResponse()
assert (False, None) == ws.can_prepare(req)
def test_closed_after_ctor():
ws = WebSocketResponse()
assert not ws.closed
assert ws.close_code is None
async def test_nonstarted_receive_json():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.receive_json()
async def test_prepare_without_upgrade(make_request):
req = make_request('GET', '/', headers=CIMultiDict({}))
ws = WebSocketResponse()
with pytest.raises(HTTPBadRequest):
await ws.prepare(req)
async def test_send_bytes_nonbytes(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
await ws.prepare(req)
with pytest.raises(TypeError):
await ws.send_bytes('string')
async def test_write_eof_not_started():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.write_eof()
async def test_nonstarted_send_str() -> None:
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.send_str('string')
async def test_nonstarted_send_json():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.send_json({'type': 'json'})
class WebsocketConnection:
def __init__(self, server, request):
self.client_id = self._get_request_client_id(request)
self.connection_id = self._get_request_connection_id(self.client_id)
self.server = server
self.request = request
self.socket = None
def _get_request_client_id(self, request):
try:
client_id = request.headers[CLIENT_ID_HEADER]
client_token = request.headers[CLIENT_TOKEN_HEADER]
except KeyError:
raise HTTPBadRequest(reason='Missing Client ID or Token Header')
with create_session() as session:
if client_id == 'None':
client = Client(token=randomstring(32))
before_client_create.emit(session=session, client=client)
session.add(client)
session.flush()
after_client_create.emit(session=session, client=client)
else:
client = session.query(Client).filter(
Client.id == client_id).first()
if not client or client.token != client_token:
raise HTTPForbidden(reason='Invalid Client ID or Token')
return client.id
def _get_request_connection_id(self, client_id):
with create_session() as session:
connection = Connection(client_id=client_id)
before_connection_create.emit(session=session,
connection=connection)
session.add(connection)
session.flush()
after_connection_create.emit(session=session,
connection=connection)
return connection.id
def _save_connection_last_activity(self):
with create_session() as session:
connection = session.query(Connection).filter(
Connection.id == self.connection_id).first()
connection.last_activity = current_timestamp()
before_connection_update.emit(session=session,
connection=connection)
session.flush()
after_connection_update.emit(session=session,
connection=connection)
def _save_connection_closed(self):
with create_session() as session:
connection = session.query(Connection).filter(
Connection.id == self.connection_id).first()
connection.closed = current_timestamp()
before_connection_update.emit(session=session,
connection=connection)
session.flush()
after_connection_update.emit(session=session,
connection=connection)
def _process_confirm(self, action):
try:
print(
'Received confirmation of action "{}" with payload: {}'.format(
action.name, action.payload),
flush=True)
action.after_confirm.emit(connection=self, action=action)
after_client_action_confirm.emit(connection=self, action=action)
except Exception as e:
print('Processing confirmation of action "{}" failed: {}'.format(
action.name, e),
flush=True)
def _process_action(self, action):
try:
print('Received action "{}" with payload: {}'.format(
action.name, action.payload),
flush=True)
before_server_action_receive.emit(connection=self, action=action)
action.before_receive.emit(connection=self, action=action)
action.execute(self)
action.after_receive.emit(connection=self, action=action)
after_server_action_receive.emit(connection=self, action=action)
self.socket.send_json(action.send_confirm())
except Exception as e:
print('Executing action "{}" failed: {}'.format(action.name, e),
flush=True)
def _process_message(self, msg):
if msg.type != WSMsgType.TEXT:
return
try:
action = parse_server_action(msg.data)
except Exception as e:
print('Invalid message received: {}; Error: {}'.format(
msg.data, e),
flush=True)
return
if action.confirm:
self._process_confirm(action)
else:
self._process_action(action)
async def connect(self):
before_server_connection_open.emit(connection=self)
self.socket = WebSocketResponse()
await self.socket.prepare(self.request)
after_server_connection_open.emit(connection=self)
self.send_client_id()
async for msg in self.socket:
self._save_connection_last_activity()
self._process_message(msg)
self._save_connection_closed()
after_server_connection_ended.emit(connection=self)
async def close(self):
before_server_connection_close.emit(connection=self)
await self.socket.close(code=WSCloseCode.GOING_AWAY)
after_server_connection_close.emit(connection=self)
def send_action(self, action):
before_client_action_send.emit(connection=self, action=action)
action.before_send.emit(connection=self, action=action)
self.socket.send_json(action.send())
action.after_send.emit(connection=self, action=action)
after_client_action_send.emit(connection=self, action=action)
async def check_send_action(self, action):
future = asyncio.Future()
sent_action = action
def listener(connection, action):
if connection is not self:
return
if action.name != sent_action.name:
return
if action.action_id != sent_action.action_id:
return
if not future.done():
future.set_result(True)
with after_client_action_confirm.connected(listener):
self.send_action(sent_action)
await asyncio.wait_for(future, ACTION_CONFIRM_TIMEOUT)
def send_client_id(self):
with create_session() as session:
client = session.query(Client).filter(
Client.id == self.client_id).first()
self.send_action(
SetIdClientAction(client_id=client.id,
client_token=client.token))
def test_can_prepare_ok(make_request):
req = make_request('GET', '/', protocols=True)
ws = WebSocketResponse(protocols=('chat', ))
assert (True, 'chat') == ws.can_prepare(req)
def test_can_prepare_unknown_protocol(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
assert (True, None) == ws.can_prepare(req)
async def test_send_json_nonjson(make_request):
req = make_request('GET', '/')
ws = WebSocketResponse()
await ws.prepare(req)
with pytest.raises(TypeError):
await ws.send_json(set())
def test_start_invalid_method(make_request):
req = make_request('POST', '/')
ws = WebSocketResponse()
with pytest.raises(HTTPMethodNotAllowed):
yield from ws.prepare(req)
async def test_write_non_prepared():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
await ws.write(b'data')
async def start_websocket_console(self, request):
"""
Connect to console using Websocket.
:param ws: Websocket object
"""
if self.status != "started":
raise NodeError("Node {} is not started".format(self.name))
if self._console_type != "telnet":
raise NodeError("Node {} console type is not telnet".format(
self.name))
try:
(telnet_reader, telnet_writer) = await asyncio.open_connection(
self._manager.port_manager.console_host, self.console)
except ConnectionError as e:
raise NodeError(
"Cannot connect to node {} telnet server: {}".format(
self.name, e))
log.info("Connected to Telnet server")
ws = WebSocketResponse()
await ws.prepare(request)
request.app['websockets'].add(ws)
log.info("New client has connected to console WebSocket")
async def ws_forward(telnet_writer):
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
telnet_writer.write(msg.data.encode())
await telnet_writer.drain()
elif msg.type == aiohttp.WSMsgType.BINARY:
await telnet_writer.write(msg.data)
await telnet_writer.drain()
elif msg.type == aiohttp.WSMsgType.ERROR:
log.debug(
"Websocket connection closed with exception {}".format(
ws.exception()))
async def telnet_forward(telnet_reader):
while not ws.closed and not telnet_reader.at_eof():
data = await telnet_reader.read(1024)
if data:
await ws.send_bytes(data)
try:
# keep forwarding websocket data in both direction
await asyncio.wait(
[ws_forward(telnet_writer),
telnet_forward(telnet_reader)],
return_when=asyncio.FIRST_COMPLETED)
finally:
log.info("Client has disconnected from console WebSocket")
if not ws.closed:
await ws.close()
request.app['websockets'].discard(ws)
return ws
def test_start_without_upgrade(make_request):
req = make_request('GET', '/',
headers=CIMultiDict({}))
ws = WebSocketResponse()
with pytest.raises(HTTPBadRequest):
yield from ws.prepare(req)
def test_write_eof_not_started():
ws = WebSocketResponse()
with pytest.raises(RuntimeError):
yield from ws.write_eof()