async def setup_two_consumers(self):
room = Room(room_name='test', password='******')
await sync_to_async(room.save)()
communicator1 = WebsocketCommunicator(application, "/ws/room/test/")
communicator2 = WebsocketCommunicator(application, "/ws/room/test/")
await communicator1.connect()
await communicator2.connect()
communicators = [communicator1, communicator2]
for idx, communicator in enumerate(communicators):
await communicator.send_json_to({
'input_type': LOGIN,
'content': {
'password': '******',
'username': '******' + str(idx)
}
})
yield communicators
await sync_to_async(room.delete)()
await communicator1.disconnect()
await communicator2.disconnect()
async def test_can_connect_to_server(self, settings):
settings.CHANNEL_LAYERS = TEST_CHANNEL_LAYERS
_, access = await create_user( # new
'*****@*****.**', 'pAssw0rd')
communicator = WebsocketCommunicator(
application=application,
path=f'/pizza/?token={access}' # changed
)
connected, _ = await communicator.connect()
assert connected is True
await communicator.disconnect()
async def test_connection_failed_for_invalid_password(self):
"""Test that clients cannot connect with an invalid password."""
# Arrange
password = PROCESS_CONNECTION_PASS + "_fake"
url = "manager/ws/subscription/?password={}".format(password)
communicator = WebsocketCommunicator(application, url)
# Act
connected, subprotocol = await communicator.connect()
# Assert
assert not connected, "Communicator should not have connected"
await communicator.disconnect()
async def test_connection_failed_for_invalid_token(self):
"""Test that clients cannot connect with an invalid token."""
# Arrange
url = "manager/ws/subscription/?token={}".format(
str(self.token) + "fake")
communicator = WebsocketCommunicator(application, url)
# Act
connected, subprotocol = await communicator.connect()
# Assert
assert not connected, "Communicator should not have connected"
await communicator.disconnect()
async def test_connection_with_password(self):
"""Test that clients can connect with a valid password."""
# Arrange
password = PROCESS_CONNECTION_PASS
url = "manager/ws/subscription/?password={}".format(password)
communicator = WebsocketCommunicator(application, url)
# Act
connected, subprotocol = await communicator.connect()
# Assert
assert connected, "Communicator was not connected"
await communicator.disconnect()
async def test_remote_reload_staggered():
communicator = WebsocketCommunicator(application, "/ws/world/sample/")
await communicator.connect()
try:
await sync_to_async(call_command)(
"connections", "force_reload", "--interval", "20", "*"
)
assert ["connection.reload", {}] == await communicator.receive_json_from()
finally:
await communicator.disconnect()
async def test_consumer_connection_init():
communicator = WebsocketCommunicator(GraphqlSubscriptionConsumer,
"/graphql/")
connected, subprotocol = await communicator.connect()
assert connected
await communicator.send_json_to({"type": "connection_init"})
response = await communicator.receive_json_from()
assert response["type"] == "connection_ack"
async def test_receive_message_for_subscribed_group_only(self):
"""Test that clients subscribed to some groups only receive messages from those."""
# Arrange
communicator = WebsocketCommunicator(application, self.url)
connected, subprotocol = await communicator.connect()
subscription_msg = {
"option": "subscribe",
"category": "telemetry",
"csc": "ScriptQueue",
"salindex": 1,
"stream": "stream1",
}
await communicator.send_json_to(subscription_msg)
await communicator.receive_json_from()
# Act
for combination in self.combinations:
msg, expected = self.build_messages(
combination["category"],
combination["csc"],
combination["salindex"],
[combination["stream"]],
)
await communicator.send_json_to(msg)
if (
combination["category"] == subscription_msg["category"]
and combination["csc"] == subscription_msg["csc"]
and combination["salindex"] == subscription_msg["salindex"]
and combination["stream"] == subscription_msg["stream"]
):
response = await communicator.receive_json_from()
# Assert
assert response == expected
else:
# Assert
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(
communicator.receive_json_from(),
timeout=self.no_reception_timeout,
)
await communicator.disconnect()
async def test_receive_message_for_subscribed_groups_only(self):
"""Test that clients subscribed to some groups only receive messages from those."""
# Arrange
communicator = WebsocketCommunicator(application, self.url)
connected, subprotocol = await communicator.connect()
for combination in self.combinations:
# Arrange: Subscribe to 1
subscription_msg = {
"csc": combination["csc"],
"stream": combination["stream"],
"category": combination["category"],
"salindex": combination["salindex"]
}
subscription_msg['option'] = 'subscribe'
await communicator.send_json_to(subscription_msg)
await communicator.receive_json_from()
# Act: Send and receive all
for combination in self.combinations:
msg, expected = \
self.build_messages(combination['category'], combination['csc'], combination['salindex'], [
combination['stream']])
await communicator.send_json_to(msg)
if combination['category'] == subscription_msg['category'] and \
combination['csc'] == subscription_msg['csc'] and \
combination['salindex'] == subscription_msg['salindex'] and \
combination['stream'] == subscription_msg['stream']:
response = await communicator.receive_json_from()
# Assert: receive the one subscribed to
assert response == expected
else:
# Assert: not receive all the others
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(
communicator.receive_json_from(),
timeout=self.no_reception_timeout)
# Clean: Unsubscribe from 1
subscription_msg['option'] = 'unsubscribe'
await communicator.send_json_to(subscription_msg)
await communicator.receive_json_from()
await communicator.disconnect()
async def test_silent_disconnection(self, admin_user, admin_client):
device = await database_sync_to_async(self._create_device)()
session_id = admin_client.cookies['sessionid'].value
communicator = WebsocketCommunicator(
self.application,
path=f'ws/controller/device/{device.pk}/',
headers=[(
b'cookie',
f'sessionid={session_id}'.encode('ascii'),
)],
)
await communicator.disconnect()
async def test_can_connect_to_server(self, settings):
settings.CHANNEL_LAYERS = TEST_CHANNEL_LAYERS
_, access = await User.objects.get(username='******')
communicator = WebsocketCommunicator(
application=application,
path='/rides/'
)
connected, _ = await communicator.connect()
assert connected is True
await communicator.disconnect()
async def test_socket_connection():
communicator = WebsocketCommunicator(
application=application,
path='/online/',
headers=[(
b'cookie',
f'sessionid=test'.encode('ascii')
)]
)
connected, _ = await communicator.connect()
assert connected is True
await communicator.disconnect()
async def test_notification(self):
# Test that parsing a bad request works
client = WebsocketCommunicator(application, 'ws/')
await client.connect()
await client.send_json_to({
"jsonrpc": "2.0",
"method": "a_notif",
"params": {}
})
assert await client.receive_nothing() is True
await client.disconnect()
async def test_push_notification_consumer__unsubscribe(user_factory):
user = await database_sync_to_async(user_factory)()
communicator = WebsocketCommunicator(PushNotificationConsumer, "/ws/notifications/")
communicator.scope["user"] = user
connected, _ = await communicator.connect()
assert connected
await communicator.send_json_to(
{"model": "user", "id": str(user.id), "action": "SUBSCRIBE"}
)
response = await communicator.receive_json_from()
assert "ok" in response
await communicator.send_json_to(
{"model": "user", "id": str(user.id), "action": "UNSUBSCRIBE"}
)
response = await communicator.receive_json_from()
assert "ok" in response
await communicator.disconnect()
async def test_bad_message():
communicator = WebsocketCommunicator(UIConsumer, "")
connected, subprotocol = await communicator.connect()
assert connected
# Test sending text
await communicator.send_json_to({"hello": "hi"})
response = await communicator.receive_json_from()
assert response['type'] == "Bad Message"
await communicator.send_json_to(
{"type": "random string with no defined response"})
response = await communicator.receive_json_from()
assert response['type'] == "Bad Message"
async def test_response_are_well_formatted(self):
# Answer should always json-rpc2
communicator = WebsocketCommunicator(application, 'ws/')
await communicator.connect()
await communicator.send_json_to({'value': 'my_value'})
response = await communicator.receive_json_from()
assert response['error']['code'] == JsonRpcConsumerTest.INVALID_REQUEST
assert response['error']['message'] == JsonRpcConsumerTest.errors[JsonRpcConsumerTest.INVALID_REQUEST]
assert response['jsonrpc'] == '2.0'
assert hasattr(response, 'id') is False
await communicator.disconnect()
async def test_receive(self):
user = mixer.blend(User, is_superuser=True)
token = Token.objects.create(user=user)
headers = [(b'Authorization', "Token {}".format(token.key).encode())]
communicator = WebsocketCommunicator(application,
'/ws/rooms/',
headers=headers)
try:
await communicator.receive_json_from()
except Exception as e:
pass
await communicator.disconnect()
async def test_can_send_message_to_websocket():
communicator = WebsocketCommunicator(ChangeListConsumer, 'ws/list')
await communicator.connect()
await communicator.send_to(text_data=json.dumps({
'type': 'admin_event',
'product': 'product',
'client': 'client',
'value': 'value'
}))
message = await communicator.receive_from()
for field in ('product', 'client', 'value'):
assert field in message
async def generate_ws_comunicators(fut, tokens, room: int, consumer, route: str):
"""
...
"""
communicators = []
for token in tokens:
communicator = WebsocketCommunicator(consumer, route)
connected, _ = await communicator.connect()
assert connected
communicators.append(communicator)
fut.set_result(communicators)
async def test_notification_subscribe_no_user():
communicator = WebsocketCommunicator(asgi.application, "/ws/notify/")
try:
connected, subprotocol = await communicator.connect()
# websocket should initially accept the connection
assert not connected
# then as there is no user, and connection must be accepted to verify, disconnect
assert await communicator.receive_nothing() is True
except Exception as e:
raise AssertionError() from e
finally:
await communicator.disconnect()
async def test_no_host_in_headers():
settings.CHANNEL_LAYERS = TEST_CHANNEL_LAYERS
room = Room.objects.create()
user = get_user_model().objects.create(username="******")
room.members.add(user)
room.save()
client = Client()
client.force_login(user=user)
with pytest.raises(ValueError):
communicator = WebsocketCommunicator(
multitenant_application, f"/ws/django_chatter/chatrooms/{room.id}/",
)
async def test_error_on_notification_frame(self):
@MyJsonRpcWebsocketConsumerTest.rpc_method()
def ping():
return True
client = WebsocketCommunicator(application, 'ws/')
await client.connect()
# we send a notification to the server
await client.send_json_to({"jsonrpc": "2.0", "method": "dwqwdq", "params": []})
self.assertEqual(await client.receive_nothing(), True)
await client.disconnect()
def test_task_flatten(self, event_loop, client, upload_obj_validated,
mocked_request):
prefix = "/uploads/" if not settings.API_PREFIX else f"/{settings.API_PREFIX}uploads/"
_, flatten_id = create_flatten(client,
upload_obj_validated,
prefix=prefix)
application = URLRouter([
re_path(r"ws/api/(?P<upload_id>[0-9a-f-]+)/$",
ValidationConsumer.as_asgi())
])
communicator = WebsocketCommunicator(
application, f"/ws/api/{upload_obj_validated.id}/")
event_loop.run_until_complete(communicator.connect())
flatten_data(flatten_id, model="Upload")
message = event_loop.run_until_complete(communicator.receive_from())
event_loop.run_until_complete(communicator.disconnect())
message = json.loads(message)
assert message["type"] == "task.flatten"
assert message["flatten"]["id"] == flatten_id
assert message["flatten"]["status"] == "processing"
async def test_join_driver_pool(self, settings):
settings.CHANNEL_LAYERS = TEST_CHANNEL_LAYERS
_, access = await create_user('test@example', 'pAss0eod', 'driver')
communicator = WebsocketCommunicator(application=application,
path=f'/taxi/?token={access}')
connected, _ = await communicator.connect()
message = {'type': 'echo.message', 'data': 'This is a test message'}
channel_layer = get_channel_layer()
await channel_layer.group_send('drivers', message=message)
response = await communicator.receive_json_from()
assert response == message
await communicator.disconnect()
async def test_auth_with_valid_token(admin_user):
communicator = WebsocketCommunicator(StreamConsumer, "/ws/streams/")
connected, subprotocol = await communicator.connect()
assert connected
await communicator.send_json_to(
data={"cmd": "auth", "token": admin_user.auth_token.key}
)
response = await communicator.receive_json_from()
assert response["code"] == 200
assert communicator.scope["user"] == admin_user
# Close
await communicator.disconnect()
async def test_auth_with_bad_token():
communicator = WebsocketCommunicator(StreamConsumer, "/ws/streams/")
connected, subprotocol = await communicator.connect()
assert connected
await communicator.send_json_to(data={
"cmd": "auth",
"api_key": "bad-auth-token"
})
response = await communicator.receive_json_from()
assert response["code"] == 400
# Close
await communicator.disconnect()
async def test_push_notification_consumer__subscribe_scratch_org_staff(
user_factory, scratch_org_factory):
user = await generate_model(user_factory, is_staff=True)
scratch_org = await generate_model(scratch_org_factory,
enqueued_at=timezone.now())
communicator = WebsocketCommunicator(PushNotificationConsumer.as_asgi(),
"/ws/notifications/")
communicator.scope["user"] = user
communicator.scope["session"] = Session()
connected, _ = await communicator.connect()
assert connected
await communicator.send_json_to({
"model": "scratchorg",
"id": str(scratch_org.id)
})
response = await communicator.receive_json_from()
assert "ok" in response
await communicator.disconnect()
async def test_consumer_delete_room():
"""
each elimination triggers signal
that is sent to all active consumers
"""
start_rooms: int = await async_count_db(Room)
communicator = WebsocketCommunicator(RoomConsumer, '/ws/rooms/')
connected, _ = await communicator.connect()
assert connected
# Test sending json
request = {
'method': 'D',
'values': {'pk_list': [3, 4]},
'token': '20fd382ed9407b31e1d5f928b5574bb4bffe6120',
}
# deleted_signal_1 = await create_event_message(
# id=3,
# operation='D',
# )
# deleted_signal_2 = await create_event_message(
# id=4,
# operation='D',
# )
# 2 signals and the answer
await communicator.send_json_to(request)
# await communicator.receive_json_from()
# signal_1 = await communicator.receive_json_from()
# assert signal_1 == deleted_signal_1 or signal_1 == deleted_signal_2
# await communicator.receive_json_from()
# signal_2 = await communicator.receive_json_from()
# assert signal_2 == deleted_signal_1 or signal_2 == deleted_signal_2
# await communicator.receive_json_from()
response = await communicator.receive_json_from()
# assert response == 'yeah'
# print(response)
assert response == {
'method': 'D',
'data': {
'count': 2,
'pk_list': [3, 4],
},
}
assert start_rooms - 2 == await async_count_db(Room)
# Close
await communicator.disconnect()
async def test_push_notification_consumer__subscribe_job(user_factory, job_factory):
user = user_factory()
job = job_factory(user=user, status=Job.Status.complete)
communicator = WebsocketCommunicator(PushNotificationConsumer, "/ws/notifications/")
communicator.scope["user"] = user
connected, _ = await communicator.connect()
assert connected
await communicator.send_json_to({"model": "job", "id": str(job.id)})
response = await communicator.receive_json_from()
assert "ok" in response
await notify_post_job(job)
response = await communicator.receive_json_from()
assert response == {
"type": "JOB_COMPLETED",
"payload": JobSerializer(instance=job, context=user_context(user)).data,
}
await communicator.disconnect()
async def test_can_send_and_receive_messages(self, settings):
settings.CHANNEL_LAYERS = TEST_CHANNEL_LAYERS
communicator = WebsocketCommunicator(application=application, path="/taxi/")
connected, _ = await communicator.connect()
message = {
"type": "echo.message",
"data": "This is a test message.",
}
await communicator.send_json_to(message)
response = await communicator.receive_json_from()
assert response == message
await communicator.disconnect()
