async def main_ws(request):
ws = WebSocketResponse()
session = await get_session(request)
expires = session.get('expires', 0)
await ws.prepare(request)
if expires < time():
await ws.close(code=4403)
return ws
user = '******'.format(
**session['user']).strip(' ')
logger.info('ws connection from %s', user)
json_str = await request.app['pg'].fetchval(calls_sql)
await ws.send_str(json_str or '[]')
request.app['ws_propagator'].add_ws(ws)
try:
async for msg in ws:
logger.info('ws message:', msg)
if msg.tp == WSMsgType.ERROR:
logger.warning('ws connection closed with exception %s',
ws.exception())
except CancelledError:
pass
finally:
request.app['ws_propagator'].remove_ws(ws)
logger.info('websocket disconnected: %s', user)
return ws
async def websocket_handler(request):
"""
Handle websocket connections
"""
ws = WebSocketResponse()
await ws.prepare(request)
logger.info("New websocket connection")
request.app["websockets"].add(ws)
snapshot_processor = request.app["snapshot_processor"]
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
try:
message = msg.json()
except (TypeError, ValueError):
await ws.send_json({"success": False, "message": "Can't load provided JSON"})
else:
await snapshot_processor.process_user_message(ws, message)
elif msg.type == aiohttp.WSMsgType.ERROR:
logger.info('ws connection closed with exception %s' % ws.exception())
logger.info("Websocket disconnect. Cleaning")
await snapshot_processor.unsubscribe(ws)
request.app["websockets"].remove(ws)
return ws
async def test_close_connections(self, mocker, bike_connection_manager):
"""Assert that closing the connection is handled gracefully."""
mocked_close = mocker.patch('aiohttp.web_ws.WebSocketResponse.close')
mocked_close.side_effect = self.conn_generator()
r0, r1 = WebSocketResponse(), WebSocketResponse()
bike_connection_manager._bike_connections[0] = r0
bike_connection_manager._bike_connections[1] = r1
await bike_connection_manager.close_connections()
assert mocked_close.call_count == 2
assert not bike_connection_manager._bike_connections
async def test_add_connection_closes_previous(self, mocker,
bike_connection_manager,
random_bike):
"""Assert that adding a new socket closes the first."""
mocked_close = mocker.patch('aiohttp.web_ws.WebSocketResponse.close')
mocked_close.side_effect = self.conn_generator()
r0, r1 = WebSocketResponse(), WebSocketResponse()
await bike_connection_manager.add_connection(random_bike, r0)
await bike_connection_manager.add_connection(random_bike, r1)
assert mocked_close.call_count == 1
async def test_send_command(self, mocker, bike_connection_manager,
random_bike):
"""Assert that you can send a connection to the bike using send_command."""
async def assert_data(json):
assert json["jsonrpc"] == "2.0"
assert json["id"] == 1
assert json["method"] == "test_command"
assert "params" in json
patched_send = mocker.patch(
'aiohttp.web_ws.WebSocketResponse.send_json')
patched_send.side_effect = assert_data
patched_counter = mocker.patch(
'server.service.manager.bike_connection_manager.BikeConnectionManager._next_rpc_id',
new_callable=PropertyMock)
patched_counter.return_value = 1
r0 = WebSocketResponse()
await bike_connection_manager.add_connection(random_bike, r0)
with pytest.raises(TimeoutError):
await bike_connection_manager._send_command(
random_bike, "test_command", timeout=timedelta(seconds=0.01))
assert patched_counter.call_count == 1
def stream(request):
redis_pool = yield from create_redis_pool(1)
subscriber = yield from redis_pool.start_subscribe()
yield from subscriber.subscribe([CURRENT_COST_KEY])
ws = WebSocketResponse()
ws.start(request)
continue_loop = True
while continue_loop:
reply = yield from subscriber.next_published()
if ws.closed:
logger.info('leaving web socket stream, usubscribing')
yield from subscriber.unsubscribe()
continue_loop = False
else:
ws.send_str(reply.value)
return ws
async def test_is_connected(self, bike_connection_manager, random_bike):
"""Assert that when the socket disappears, so does the bike connection."""
r0 = WebSocketResponse()
await bike_connection_manager.add_connection(random_bike, r0)
bike_connection_manager._bike_locations[random_bike.id] = None
bike_connection_manager._bike_battery[random_bike.id] = None
bike_connection_manager._bike_locked[random_bike.id] = None
assert bike_connection_manager.is_connected(random_bike)
del r0
assert not bike_connection_manager.is_connected(random_bike)
async def listen_for_close(self, ws: WebSocketResponse) -> None:
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
if msg.data == 'close':
await ws.close()
logger.info("closed websocket")
if msg.data == 'pong':
self.pong_event.set()
elif msg.type == aiohttp.WSMsgType.ERROR:
logger.error('ws connection closed with exception %s' %
ws.exception())
async def test_add_connection_closed(self, mocker, bike_connection_manager,
random_bike):
"""Assert that adding a closed connection fails."""
mocked_close = mocker.patch('aiohttp.web_ws.WebSocketResponse.close')
mocked_close.side_effect = self.conn_generator()
mocked_closed = mocker.patch('aiohttp.web_ws.WebSocketResponse.closed')
mocked_closed.return_value = True
with pytest.raises(ConnectionError):
await bike_connection_manager.add_connection(
random_bike, WebSocketResponse())
class WebsocketHandler:
def __init__(self, request):
self.request = request
self.remote = request.remote
self.peername = request.transport.get_extra_info('peername')
self.pubsub = request.app['pubsub']
self.response = WebSocketResponse()
def __repr__(self):
return f"<WebsocketSubscriber {self.remote}: {self.peername}>"
async def send_metric(self, key, value):
message = json_dumps({key: value})
await self.response.send_str(message)
async def send_json(self, *a, **kw):
data = dict(*a, **kw)
message = json_dumps(data)
await self.response.send_str(message)
async def loop(self):
await self.response.prepare(self.request)
logger.debug("ws prepared")
async for msg in self.response:
if msg.type == WSMsgType.TEXT:
try:
data = json.loads(msg.data)
self._parse_command(data)
except Exception as e:
await self.send_json(error=format_exception(e))
logger.exception('Exception in message %r', msg)
elif msg.type == WSMsgType.ERROR:
logger.error('ws connection closed with exception %s',
self.response.exception())
break
else:
logger.warning("ws message: %s", msg)
self.pubsub.unsubscribe(self)
logger.debug("Websocket closed: %r", self)
return self.response
def _parse_command(self, data):
if 'subscribe' in data and isinstance(data['subscribe'], str):
self.pubsub.subscribe(data['subscribe'], self)
else:
raise ValueError("invalid command")
async def handle(self, request):
ws = WebSocketResponse()
await ws.prepare(request)
handler = GrozaHandler()
conn = GrozaServerConnection(handler,
WsAiohttpServerTransportResponse(ws))
await self._server.notify_connection_start(conn)
try:
await conn.handle()
finally:
await self._server.notify_connection_close(conn)
return ws
async def news_websocket_handler(request):
ws = WebSocketResponse()
await ws.prepare(request)
profile_id = 0
async for ws_msg in ws:
if ws_msg.type == aiohttp.WSMsgType.TEXT:
profile_id = await check_has_profile(ws_msg.data, request['conn'])
break
queue = request.app['queue']
async for msg in queue.receive():
if msg['profile']['id'] in await followed_profiles(request['conn'], profile_id):
name = f'{msg["profile"]["first_name"]} {msg["profile"]["last_name"]}'
await ws.send_str(f'''
<p>
<a href="/profiles/post/{msg["id"]}">{msg["text"][:50]}...</a> by
<a href="{msg["profile"]["id"]}">{name}</a></p>''')
return ws
async def test_resolve_command(self, mocker, bike_connection_manager,
random_bike):
"""Assert that an opened request can be resolved."""
async def noop(json):
pass
patched_send = mocker.patch(
'aiohttp.web_ws.WebSocketResponse.send_json')
patched_send.side_effect = noop
r0 = WebSocketResponse()
rpc = RPC(1, r0, "command")
bike_connection_manager._pending_commands[random_bike.id][1] = rpc
rpc_request = rpc()
rpc_resolver = bike_connection_manager.resolve_command(
random_bike, 1, "returned_data")
returned_data, _ = await gather(rpc_request, rpc_resolver)
assert returned_data == "returned_data"
async def websocket_handler(request):
resp = WebSocketResponse()
await resp.prepare(request)
async with AsyncTail(filepath=request.app['filepath']) as atail:
async for line in atail:
resp.send_str(json.dumps({'action': 'sent', 'text': line}))
resp.send_str(json.dumps({
'action': 'close',
}))
await resp.close()
return resp
def __init__(self, request):
self.request = request
self.remote = request.remote
self.peername = request.transport.get_extra_info('peername')
self.pubsub = request.app['pubsub']
self.response = WebSocketResponse()
async def unload_sock(request: Request):
"""
WebSocket view to detect when Messenger closes the WebView.
There is a dual mechanism:
- If "unload" is received over the socket, then close instantly
- If no heartbeat is received for some time, them close
The URL must include signed information about the user, like what
the `sign_webview` parameter of UrlButton would provide.
"""
tk = request.query.get(settings.WEBVIEW_TOKEN_KEY)
if not tk:
return json_response(
{
'error': True,
'message': 'Missing "{}"'.format(settings.WEBVIEW_TOKEN_KEY),
},
status=400)
try:
tk = jwt.decode(tk, settings.WEBVIEW_SECRET_KEY)
except jwt.InvalidTokenError:
return json_response(
{
'error': True,
'message': 'Provided token is invalid'
},
status=400)
try:
user_id = tk['fb_psid']
assert isinstance(user_id, Text)
page_id = tk['fb_pid']
assert isinstance(page_id, Text)
slug = tk['slug']
assert isinstance(slug, Text) or slug is None
except (KeyError, AssertionError):
return json_response(
{
'error': True,
'message': 'Provided payload is invalid'
},
status=400)
event = {
'sender': {
'id': user_id,
},
'recipient': {
'id': page_id,
},
'url_base': str(request.url),
'close_webview': {
'slug': slug,
},
}
ws = WebSocketResponse(timeout=settings.WEBVIEW_HEARTBEAT_TIMEOUT,
receive_timeout=settings.WEBVIEW_HEARTBEAT_TIMEOUT)
await ws.prepare(request)
inhibit = False
try:
# noinspection PyTypeChecker
async for msg in ws:
# For some reason, the timeout does not work without a print
# Yeah, what the f**k right?
# TODO make it work properly
print('hack print')
if msg.type == WSMsgType.TEXT:
if msg.data == 'unload':
break
if msg.data == 'inhibit':
inhibit = True
elif msg.type == WSMsgType.ERROR:
break
except asyncio.TimeoutError:
pass
if not inhibit:
fb = await manager.get_platform('facebook')
msg = FacebookMessage(event, fb, False)
await fb.handle_event(msg)
return ws
async def test_add_connection(self, bike_connection_manager, random_bike):
"""Assert that the list of connected bikes is maintained."""
r0 = WebSocketResponse()
await bike_connection_manager.add_connection(random_bike, r0)
assert len(list(
bike_connection_manager._bike_connections.values())) == 1