kokkoja-rel/tests/test_relay.py

"""Tests of the relay.  Run from the project directory:

    .venv/bin/python -m unittest discover tests
"""

import asyncio
import json
import pathlib
import sys
import unittest

from websockets.asyncio.client import connect
from websockets.asyncio.server import serve
from websockets.exceptions import ConnectionClosed, InvalidStatus

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent / "relay"))
import kokkoja_relay as relay  # noqa: E402

ROOM_A = "A" * 22
ROOM_B = "B" * 22


def to_peer(destination, body):
    return destination.to_bytes(4, "big") + body


class RelayTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        relay.rooms.clear()
        self.server = await serve(relay.handle, "127.0.0.1", 0,
                                  process_request=relay.check_path,
                                  max_size=relay.MAX_MESSAGE_BYTES)
        port = self.server.sockets[0].getsockname()[1]
        self.address = f"ws://127.0.0.1:{port}"

    async def asyncTearDown(self):
        self.server.close()
        await self.server.wait_closed()

    async def join(self, room):
        connection = await connect(f"{self.address}/room/{room}")
        self.addAsyncCleanup(connection.close)
        welcome = json.loads(await connection.recv())
        self.assertEqual(welcome["type"], "welcome")
        return connection, welcome

    async def test_welcome_join_leave(self):
        first, welcome = await self.join(ROOM_A)
        self.assertEqual(welcome, {"type": "welcome", "peer": 1, "peers": []})
        second, welcome = await self.join(ROOM_A)
        self.assertEqual(welcome, {"type": "welcome", "peer": 2, "peers": [1]})
        self.assertEqual(json.loads(await first.recv()), {"type": "join", "peer": 2})
        await second.close()
        self.assertEqual(json.loads(await first.recv()), {"type": "leave", "peer": 2})

    async def test_broadcast_reaches_the_others_only(self):
        first, _ = await self.join(ROOM_A)
        second, _ = await self.join(ROOM_A)
        third, _ = await self.join(ROOM_A)
        await first.recv()      # join 2
        await first.recv()      # join 3
        await second.recv()     # join 3
        await second.send(to_peer(0, b"hello"))
        self.assertEqual(await first.recv(), to_peer(2, b"hello"))
        self.assertEqual(await third.recv(), to_peer(2, b"hello"))
        # The sender gets nothing back: the next thing it receives is this.
        await first.send(to_peer(0, b"next"))
        self.assertEqual(await second.recv(), to_peer(1, b"next"))

    async def test_directed_message_reaches_one_peer(self):
        first, _ = await self.join(ROOM_A)
        second, _ = await self.join(ROOM_A)
        third, _ = await self.join(ROOM_A)
        await first.recv()
        await first.recv()
        await second.recv()
        await first.send(to_peer(3, b"for three"))
        await first.send(to_peer(0, b"for all"))
        self.assertEqual(await third.recv(), to_peer(1, b"for three"))
        self.assertEqual(await third.recv(), to_peer(1, b"for all"))
        self.assertEqual(await second.recv(), to_peer(1, b"for all"))

    async def test_rooms_are_separate(self):
        in_a, _ = await self.join(ROOM_A)
        in_b, welcome = await self.join(ROOM_B)
        self.assertEqual(welcome["peers"], [])
        other_a, _ = await self.join(ROOM_A)
        await in_a.recv()       # join
        await in_b.send(to_peer(0, b"b only"))
        await other_a.send(to_peer(0, b"a only"))
        self.assertEqual(await in_a.recv(), to_peer(2, b"a only"))

    async def test_room_is_forgotten_when_empty(self):
        first, _ = await self.join(ROOM_A)
        await first.close()
        await first.wait_closed()
        for _ in range(50):
            if not relay.rooms:
                break
            await asyncio.sleep(0.01)
        self.assertEqual(relay.rooms, {})

    async def test_bad_path_is_refused(self):
        for path in ("/", "/room/short", f"/room/{ROOM_A}/extra", "/room/" + "!" * 22):
            with self.assertRaises(InvalidStatus):
                await connect(self.address + path)

    async def test_text_message_closes_the_connection(self):
        first, _ = await self.join(ROOM_A)
        await first.send("text is not allowed")
        with self.assertRaises(ConnectionClosed) as caught:
            await first.recv()
        self.assertEqual(caught.exception.rcvd.code, relay.CLOSE_BAD_MESSAGE)

    async def test_room_full(self):
        relay.MAX_PEERS_PER_ROOM, saved = 2, relay.MAX_PEERS_PER_ROOM
        try:
            await self.join(ROOM_A)
            await self.join(ROOM_A)
            extra = await connect(f"{self.address}/room/{ROOM_A}")
            with self.assertRaises(ConnectionClosed) as caught:
                await extra.recv()
            self.assertEqual(caught.exception.rcvd.code, relay.CLOSE_ROOM_FULL)
        finally:
            relay.MAX_PEERS_PER_ROOM = saved


if __name__ == "__main__":
    unittest.main()