def test_create_channel(self):
message = self.sample_message
message.type = "create_channel"
message.receiver_guid = ""
message.data = {"name": "new_channel"}
# Mock the webocket methods
self.sender_mock_socket.send = unittest.mock.AsyncMock()
dispatcher = Dispatcher({})
self.mock_receive_message(message, dispatcher)
# It adds a new channel
found = False
new_channel = ""
for channel in dispatcher.channels.values():
found = channel.name == "new_channel"
if found:
new_channel = channel
break
self.assertTrue(found)
# It sends an 'ack' with the channel id
self.sender_mock_socket.send.assert_called()
sent_message = Message.from_json(
self.sender_mock_socket.send.call_args_list[0][0][0])
self.assertEqual(sent_message.type, "ack")
self.assertEqual(sent_message.data["channel_id"], new_channel.id)
def test_send(self, SenderMockWebSocket, ReceiverMockWebSocket):
# Set up some clients
sender = Client(SenderMockWebSocket)
receiver = Client(ReceiverMockWebSocket)
message_handler = MessageHandler()
message_handler.add_client(sender)
message_handler.add_client(receiver)
self.assertTrue(len(message_handler.clients) == 2)
# It sends a message to receiver
sample_message = Message(receiver_guid=receiver.address,
type="announce",
data="data")
asyncio.get_event_loop().run_until_complete(
message_handler.send(sender, sample_message))
# It called send
ReceiverMockWebSocket.send.assert_called()
# It sent the right message
sent_message = Message.from_json(
ReceiverMockWebSocket.send.call_args[0][0])
sent_data = sent_message.data
sent_type = sent_message.type
self.assertEqual(sent_data, sample_message.data)
self.assertEqual(sent_type, sample_message.type)
# It set the "sender_guid"
sent_from = sent_message.sender_guid
self.assertEqual(sent_from, sender.address)
def test_broadcast(self, SenderMockWebSocket, Receiver1MockWebSocket, Receiver2MockWebSocket):
# Set up some clients
sender = Client(SenderMockWebSocket)
receiver1 = Client(Receiver1MockWebSocket)
receiver2 = Client(Receiver2MockWebSocket)
channel = Channel("test_channel")
channel.add_client(sender)
channel.add_client(receiver1)
channel.add_client(receiver2)
self.assertTrue(len(channel.clients) == 3)
# It sends a message to each client
sample_message = Message(
type = "announce",
data= "data"
)
asyncio.get_event_loop().run_until_complete(
channel.broadcast(sender, sample_message)
)
# It called send
Receiver1MockWebSocket.send.assert_called()
Receiver2MockWebSocket.send.assert_called()
# It sent the right message
sent_message = Message.from_json(Receiver1MockWebSocket.send.call_args[0][0])
sent_data = sent_message.data
sent_type = sent_message.type
sent_channel_id = sent_message.channel_id
self.assertEqual(sent_data, sample_message.data)
self.assertEqual(sent_type, sample_message.type)
self.assertEqual(sent_channel_id, channel.id)
# It set the "sender_guid"
sent_from = sent_message.sender_guid
self.assertEqual(sent_from, sender.address)
async def dispatch(websocket, path):
client = Client(websocket)
print("New client: " + client.address, flush=True)
try:
async for json_message in websocket:
message = Message.from_json(json_message)
if (message.type == "announce"):
message_handler.add_client(client)
await message_handler.broadcast(client, message)
elif (message.type in ["offer","answer","ice"]):
await message_handler.send(client, message)
finally:
message_handler.remove_client(client)
# broadcast the departure
message = Message(type = "hangup")
await message_handler.broadcast(client, message)
print("Client " + client.address + " left", flush=True)
async def handle_connection(self, websocket: iter, path: str):
client = Client(websocket)
print("New client: " + client.address, flush=True)
channel = None
try:
async for json_message in websocket:
message = Message.from_json(json_message)
if message.type == "create_channel":
new_channel = Channel(message.data["name"])
self.channels[new_channel.id] = new_channel
ack = Message(type="ack",
data={"channel_id": new_channel.id})
await websocket.send(ack.to_json())
continue
channel_id = message.channel_id
if channel_id in self.channels.keys():
channel = self.channels[channel_id]
else:
message = Message(type="error")
await websocket.send(message.to_json())
await websocket.close()
break
if (message.type in ["announce"]):
channel.add_client(client)
await channel.broadcast(client, message)
elif (message.type in ["offer", "answer", "ice"]):
await channel.send(client, message)
elif (message.type in ["name"]):
await channel.broadcast(client, message)
finally:
if channel:
channel.remove_client(client)
# broadcast the departure
message = Message(type="hangup")
await channel.broadcast(client, message)
print("Client " + client.address + " left", flush=True)
def test_message_with_missing_channel(self):
# When channel_id is missing
message = self.sample_message
message.type = "anything_but_create_channel"
message.channel_id = "non_existent"
message.receiver_guid = ""
# Mock the webocket methods
self.sender_mock_socket.close = unittest.mock.AsyncMock()
self.sender_mock_socket.send = unittest.mock.AsyncMock()
# Receive the test message
dispatcher = Dispatcher(self.channels)
self.mock_receive_message(message, dispatcher)
# It sends an error to the sender
self.sender_mock_socket.send.assert_called()
sent_message = Message.from_json(
self.sender_mock_socket.send.call_args_list[0][0][0])
self.assertEqual(sent_message.type, "error")
# It disconnects
self.sender_mock_socket.close.assert_called()
def test_from_json(self):
message_json = json.dumps(sample_message)
message = Message.from_json(message_json)
self.assertEqual(message.__dict__, sample_message)
