def test_mismatched_key(self):
@self.client.set_mock_api
def empty_dict_allowed(request, **kwargs):
return MockResponse(
{
'user': {
'id': 1
},
'allowed_rooms': [{
'target': 'car',
'foo': 1,
}],
}, 200)
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_car))
self.client.open_connection(ws)
self.assertEqual(2, len(ws.outgoing_messages))
self.assertEqual(
{
'requestId': 'a',
'code': 'error',
'message': 'room-not-found'
}, json.loads(ws.outgoing_messages[1]))
def test_ping_pong(self):
ws = MockWebSocket()
ws.mock_incoming_message('ping')
self.client.open_connection(ws)
self.assertEqual(2, len(ws.outgoing_messages))
self.assertEqual('pong', ws.outgoing_messages[1])
def test_ip_forward(self):
ws = MockWebSocket()
# Fake nginx IP resolving
ws.environ['HTTP_X_REAL_IP'] = '1.2.3.4'
self.client.open_connection(ws)
self.assertEqual(1, self.client.mock_send.call_count)
request = self.client.mock_send.call_args_list[0][0][0]
self.assertIn('X-Forwarded-For', request.headers)
self.assertEqual('1.2.3.4', request.headers['X-Forwarded-For'])
def test_subscribe_success(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_car))
self.client.open_connection(ws)
self.assertEqual({
'requestId': 'a',
'code': 'success'
}, json.loads(ws.outgoing_messages[1]))
def test_room_key_order_unimportant(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_car_reverse))
self.client.open_connection(ws)
self.assertEqual({
'requestId': 'a',
'code': 'success'
}, json.loads(ws.outgoing_messages[1]))
def test_subscribe_unallowed_room(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_weird))
self.client.open_connection(ws)
self.assertEqual(
{
'requestId': 'b',
'code': 'error',
'message': 'room-not-found'
}, json.loads(ws.outgoing_messages[1]))
def test_subscribe_is_rejected_for_unlisted_room(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_other_car))
hub = self.client.app.hub
self.assertEqual(0, len(hub.rooms.keys()))
g1 = greenlet(self.client.open_connection)
g1.switch(ws)
self.assertHubRoomsEqual([])
def test_cannot_unsubscribe_when_not_subscribed(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(unsubscribe_car))
self.client.open_connection(ws)
self.assertEqual(2, len(ws.outgoing_messages))
self.assertEqual(
{
'requestId': 'car',
'code': 'error',
'message': 'not-subscribed'
}, json.loads(ws.outgoing_messages[1]))
def test_close_empty_room(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_ride))
hub = self.client.app.hub
g1 = greenlet(self.client.open_connection)
g1.switch(ws)
self.assertHubRoomsEqual([room_ride])
ws.connection.unsubscribe_all()
self.assertEqual(0, len(hub.rooms.keys()))
def test_unauth_closed(self):
@self.client.set_mock_api
def not_authenticated(request, **kwargs):
return MockResponse({}, 403)
ws = MockWebSocket()
self.client.open_connection(ws)
self.assertEqual(True, ws.closed)
def test_incoming_websocket_calls_bootstrap(self):
ws = MockWebSocket()
self.client.open_connection(ws)
self.assertEqual(1, self.client.mock_send.call_count)
external_url = '{}bootstrap/'.format(api_url)
request = self.client.mock_send.call_args_list[0][0][0]
self.assertEqual('GET', request.method)
self.assertEqual(external_url, request.url)
def test_wildcard(self):
allowed_car_all = {'target': 'car', 'car': '*'}
@self.client.set_mock_api
def allowed_car_wildcard(request, **kwargs):
return MockResponse(
{
'user': {
'id': 1
},
'allowed_rooms': [allowed_car_all],
}, 200)
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_car))
self.client.open_connection(ws)
self.assertEqual({
'requestId': 'a',
'code': 'success'
}, json.loads(ws.outgoing_messages[1]))
def test_subscribe_specific_is_accepted_for_wildcard_room(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_bicycle_specific))
hub = self.client.app.hub
self.assertEqual(0, len(hub.rooms.keys()))
g1 = greenlet(self.client.open_connection)
g1.switch(ws)
self.assertHubRoomsEqual([room_bicycle_specific])
# Test that a room get a reference to its connections
# and vice versa
room = self.getHubRoomByDict(room_bicycle_specific)
connection = room.connections[0]
self.assertEqual(ws, connection.ws)
self.assertEqual([room], [s.room for s in connection.subscriptions.values()])
def test_subscribe_creates_new(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_ride))
hub = self.client.app.hub
self.assertEqual(0, len(hub.rooms.keys()))
g1 = greenlet(self.client.open_connection)
g1.switch(ws)
self.assertHubRoomsEqual([room_ride])
# Test that a room get a reference to its connections
# and vice versa
room = self.getHubRoomByDict(room_ride)
connection = room.connections[0]
self.assertEqual(ws, connection.ws)
self.assertEqual([room], [s.room for s in connection.subscriptions.values()])
def test_auth_but_not_logged_in_closed(self):
@self.client.set_mock_api
def not_logged_in(request, **kwargs):
return MockResponse({'user': None}, 200)
ws = MockWebSocket()
self.client.open_connection(ws)
# We do get a response, but the rooms are simply empty if
# there's no user object in the response.
res = json.loads(ws.outgoing_messages[0])
self.assertCountEqual([], res['allowed_rooms'])
self.assertEqual(True, ws.closed)
def test_connection_headers(self):
class WerkzeugRequest:
def __init__(self, value):
self.value = value
@property
def args(self):
return {'foo': self.value}
ws = MockWebSocket()
ws.environ['werkzeug.request'] = WerkzeugRequest('bar')
self.client.open_connection(ws)
self.assertEqual(1, self.client.mock_send.call_count)
request = self.client.mock_send.call_args_list[0][0][0]
self.assertIn('foo', request.headers)
self.assertEqual('BAR', request.headers['foo'])
def test_token_query_param_to_auth_header(self):
class WerkzeugRequest:
def __init__(self, token):
self.token = token
@property
def args(self):
return {'token': self.token}
ws = MockWebSocket()
ws.environ['werkzeug.request'] = WerkzeugRequest('itsmeyourelookingfor')
self.client.open_connection(ws)
self.assertEqual(1, self.client.mock_send.call_count)
request = self.client.mock_send.call_args_list[0][0][0]
self.assertIn('Authorization', request.headers)
self.assertEqual('Token itsmeyourelookingfor', request.headers['Authorization'])
def test_subscribe_joins_existing(self):
ws1 = MockWebSocket()
ws1.mock_incoming_message(json.dumps(subscribe_ride))
ws2 = MockWebSocket()
ws2.mock_incoming_message(json.dumps(subscribe_ride))
g1 = greenlet(self.client.open_connection)
g1.switch(ws1)
self.assertHubRoomsEqual([room_ride])
room = self.getHubRoomByDict(room_ride)
c1 = room.connections[0]
g2 = greenlet(self.client.open_connection)
g2.switch(ws2)
self.assertHubRoomsEqual([room_ride])
c2 = room.connections[1]
self.assertEqual([c1, c2], room.connections)
self.assertEqual([room], [s.room for s in c1.subscriptions.values()])
self.assertEqual([room], [s.room for s in c2.subscriptions.values()])
def test_no_more_publish(self):
ws = MockWebSocket()
ws.mock_incoming_message(json.dumps(subscribe_car))
g = greenlet(self.client.open_connection)
g.switch(ws)
self.assertEqual(2, len(ws.outgoing_messages))
self.assertEqual({
'requestId': 'car',
'code': 'success'
}, json.loads(ws.outgoing_messages[1]))
self.trigger({'rooms': [room_car], 'data': [{'id': 1}]})
g.switch()
self.assertEqual(3, len(ws.outgoing_messages))
self.assertEqual(
{
'requestId': 'car',
'type': 'publish',
'data': [{
'id': 1
}]
}, json.loads(ws.outgoing_messages[2]))
ws.mock_incoming_message(json.dumps(unsubscribe_car))
g.switch()
self.assertEqual(4, len(ws.outgoing_messages))
self.assertEqual({
'requestId': 'car',
'code': 'success'
}, json.loads(ws.outgoing_messages[3]))
self.trigger({'rooms': [room_car], 'data': [{'id': 1}]})
g.switch()
self.assertEqual(4, len(ws.outgoing_messages))
def test_keeps_open_nonempty_room(self):
ws1 = MockWebSocket()
ws1.mock_incoming_message(json.dumps(subscribe_ride))
ws2 = MockWebSocket()
ws2.mock_incoming_message(json.dumps(subscribe_ride))
ws2.mock_incoming_message(json.dumps(subscribe_car))
g1 = greenlet(self.client.open_connection)
g1.switch(ws1)
g2 = greenlet(self.client.open_connection)
g2.switch(ws2)
# Test that both rooms are created and contain
# the correct connections
self.assertHubRoomsEqual([room_ride, room_car])
r_car = self.getHubRoomByDict(room_car)
self.assertEqual([ws2.connection], r_car.connections)
r_ride = self.getHubRoomByDict(room_ride)
self.assertSetEqual(set([ws1.connection, ws2.connection]), set(r_ride.connections))
# Close the second websocket, unsubscribing if for all the rooms
ws2.close()
# Test that the now empty room_car is closed
self.assertHubRoomsEqual([room_ride])
# And test that the room_ride still contains ws1
self.assertEqual([ws1.connection], r_ride.connections)
def test_sub_removed_from_room(self):
ws1 = MockWebSocket()
ws1.mock_incoming_message(json.dumps(subscribe_car))
ws2 = MockWebSocket()
ws2.mock_incoming_message(json.dumps(subscribe_car_2))
g1 = greenlet(self.client.open_connection)
g1.switch(ws1)
g2 = greenlet(self.client.open_connection)
g2.switch(ws2)
room = self.getHubRoomByDict(room_car)
c1 = room.connections[0]
c2 = room.connections[1]
self.assertHubRoomsEqual([room_car])
self.assertEqual([c1, c2], room.connections)
self.assertEqual([room], [s.room for s in c1.subscriptions.values()])
self.assertEqual([room], [s.room for s in c2.subscriptions.values()])
ws2.mock_incoming_message(json.dumps(unsubscribe_car_2))
g2.switch()
self.assertHubRoomsEqual([room_car])
self.assertEqual([c1], room.connections)
self.assertEqual([room], [s.room for s in c1.subscriptions.values()])
self.assertEqual(0, len(c2.subscriptions))
ws1.mock_incoming_message(json.dumps(unsubscribe_car))
g1.switch()
self.assertHubRoomsEqual([])
self.assertEqual(0, len(c1.subscriptions))
self.assertEqual(0, len(c2.subscriptions))
def test_ip_forward_not_provided(self):
ws = MockWebSocket()
self.client.open_connection(ws)
self.assertEqual(1, self.client.mock_send.call_count)
request = self.client.mock_send.call_args_list[0][0][0]
self.assertNotIn('X-Forwarded-For', request.headers)
def test_incoming_trigger_gets_published_on_wildcard_rooms_for_only_wildcard_subscribers(self):
ws1 = MockWebSocket()
ws1.mock_incoming_message(json.dumps(subscribe_bicycle_specific))
ws2 = MockWebSocket()
ws2.mock_incoming_message(json.dumps(subscribe_other_bicycle_specific))
ws3 = MockWebSocket()
ws3.mock_incoming_message(json.dumps(subscribe_bicycle_wildcard))
g1 = greenlet(self.client.open_connection)
g1.switch(ws1)
g2 = greenlet(self.client.open_connection)
g2.switch(ws2)
g3 = greenlet(self.client.open_connection)
g3.switch(ws3)
# Mock incoming ride trigger
self.client.flask_test_client.post(
'/trigger/',
content_type='application/json',
data=json.dumps({
'rooms': [room_bicycle_wildcard],
'data': [{
'id': 1,
}]
}))
# The first outgoing message is the allowed_rooms listing
# The second one is about the successful subscribe
# The third is the actual publish
self.assertEqual(2, len(ws1.outgoing_messages))
self.assertEqual(2, len(ws2.outgoing_messages))
self.assertEqual(3, len(ws3.outgoing_messages))
self.assertEqual(json.dumps(publish_bicycle), ws3.outgoing_messages[2])
def test_incoming_trigger_gets_published(self):
ws1 = MockWebSocket()
ws1.mock_incoming_message(json.dumps(subscribe_ride))
ws2 = MockWebSocket()
ws2.mock_incoming_message(json.dumps(subscribe_ride))
ws2.mock_incoming_message(json.dumps(subscribe_car))
g1 = greenlet(self.client.open_connection)
g1.switch(ws1)
g2 = greenlet(self.client.open_connection)
g2.switch(ws2)
# Mock incoming ride trigger
self.client.flask_test_client.post(
'/trigger/',
content_type='application/json',
data=json.dumps({
'rooms': [room_ride],
'data': [{
'id': 1,
}]
}))
# The first outgoing message is the allowed_rooms listing
# The second one is about the successful subscribe
# The third is the actual publish
self.assertEqual(3, len(ws1.outgoing_messages))
self.assertEqual(json.dumps(publish_ride), ws1.outgoing_messages[2])
# ws2 has 2 subscribe success events
self.assertEqual(4, len(ws2.outgoing_messages))
self.assertEqual(json.dumps(publish_ride), ws2.outgoing_messages[3])
ws2.close()
self.client.flask_test_client.post(
'/trigger/',
content_type='application/json',
data=json.dumps({
'rooms': [room_ride],
'data': [{
'id': 1,
}]
}))
self.assertEqual(4, len(ws1.outgoing_messages))
self.assertEqual(json.dumps(publish_ride), ws1.outgoing_messages[3])
def test_silent_close(self):
ws1 = MockWebSocket()
ws1.mock_incoming_message(json.dumps(subscribe_ride))
ws2 = MockWebSocket()
ws2.mock_incoming_message(json.dumps(subscribe_ride))
ws2.mock_incoming_message(json.dumps(subscribe_car))
g1 = greenlet(self.client.open_connection)
g1.switch(ws1)
g2 = greenlet(self.client.open_connection)
g2.switch(ws2)
# Silent close the websocket, the exceptions still exist
ws2.closed = True
# Mock incoming ride trigger
self.client.flask_test_client.post(
'/trigger/',
content_type='application/json',
data=json.dumps({
'rooms': [room_ride],
'data': [{
'id': 1,
}]
}))
# The first outgoing message is the allowed_rooms listing
# The second one is about the successful subscribe
# The third is the actual publish
self.assertEqual(3, len(ws1.outgoing_messages))
self.assertEqual(json.dumps(publish_ride), ws1.outgoing_messages[2])
# Make sure the closed websocket has its connection
# removed from the rooms
self.assertHubRoomsEqual([room_ride])
r_ride = self.getHubRoomByDict(room_ride)
self.assertEqual([ws1.connection], r_ride.connections)
def test_lists_allowed_rooms(self):
ws = MockWebSocket()
self.client.open_connection(ws)
res = json.loads(ws.outgoing_messages[0])
self.assertCountEqual([room_ride, room_car, room_bicycle_wildcard], res['allowed_rooms'])