devplacepy/tests/api/messages/ws.py

# retoor <retoor@molodetz.nl>

import asyncio
import json
import time

import pytest
import requests
import websockets

from tests.conftest import BASE_URL, PORT
from devplacepy.database import get_table, refresh_snapshot, set_setting

WS_URL = f"ws://127.0.0.1:{PORT}/messages/ws"
_counter = [0]


@pytest.fixture(scope="module", autouse=True)
def _ws_settings(app_server):
    for key, value in {
        "rate_limit_per_minute": "1000000",
        "rate_limit_window_seconds": "60",
        "registration_open": "1",
        "maintenance_mode": "0",
    }.items():
        set_setting(key, value)
    yield


def _db_user(name):
    refresh_snapshot()
    return get_table("users").find_one(username=name)


def _unique(prefix="wsm"):
    _counter[0] += 1
    return f"{prefix}{int(time.time() * 1000)}{_counter[0]}"


def _member():
    name = _unique()
    s = requests.Session()
    s.post(
        f"{BASE_URL}/auth/signup",
        data={
            "username": name,
            "email": f"{name}@t.dev",
            "password": "secret123",
            "confirm_password": "secret123",
            "birth_date": "1990-01-01",
            "accept_terms": "1",
        },
        allow_redirects=True,
    )
    return s, name


def _cookie_header(session):
    token = session.cookies.get("session")
    assert token, "expected a session cookie after signup"
    return {"Cookie": f"session={token}"}


async def _recv_until(ws, frame_type, timeout=5.0):
    deadline = asyncio.get_event_loop().time() + timeout
    while True:
        remaining = deadline - asyncio.get_event_loop().time()
        if remaining <= 0:
            raise AssertionError(f"timed out waiting for {frame_type}")
        raw = await asyncio.wait_for(ws.recv(), timeout=remaining)
        frame = json.loads(raw)
        if frame.get("type") == frame_type:
            return frame


def test_ws_requires_authentication(app_server):
    async def run():
        async with websockets.connect(WS_URL) as ws:
            with pytest.raises(websockets.exceptions.ConnectionClosed):
                await asyncio.wait_for(ws.recv(), timeout=3.0)

    asyncio.run(run())


def test_ws_send_broadcasts_to_both_parties(app_server):
    sender_session, _ = _member()
    receiver_session, receiver_name = _member()
    receiver_uid = _db_user(receiver_name)["uid"]
    client_id = "cid-" + _unique("c")

    async def run():
        async with websockets.connect(
            WS_URL, additional_headers=_cookie_header(sender_session)
        ) as sender_ws, websockets.connect(
            WS_URL, additional_headers=_cookie_header(receiver_session)
        ) as receiver_ws:
            await _recv_until(sender_ws, "ready")
            await _recv_until(receiver_ws, "ready")

            await sender_ws.send(
                json.dumps(
                    {
                        "type": "send",
                        "receiver_uid": receiver_uid,
                        "content": "hello over the wire",
                        "client_id": client_id,
                    }
                )
            )

            echo = await _recv_until(sender_ws, "message")
            assert echo["client_id"] == client_id
            assert echo["content"] == "hello over the wire"
            assert echo["receiver_uid"] == receiver_uid

            delivered = await _recv_until(receiver_ws, "message")
            assert delivered["content"] == "hello over the wire"
            assert delivered["sender_uid"] == echo["sender_uid"]

    asyncio.run(run())


def test_ws_typing_relays_to_receiver_only(app_server):
    sender_session, sender_name = _member()
    receiver_session, receiver_name = _member()
    sender_uid = _db_user(sender_name)["uid"]
    receiver_uid = _db_user(receiver_name)["uid"]

    async def run():
        async with websockets.connect(
            WS_URL, additional_headers=_cookie_header(sender_session)
        ) as sender_ws, websockets.connect(
            WS_URL, additional_headers=_cookie_header(receiver_session)
        ) as receiver_ws:
            await _recv_until(sender_ws, "ready")
            await _recv_until(receiver_ws, "ready")

            await sender_ws.send(
                json.dumps({"type": "typing", "receiver_uid": receiver_uid})
            )
            typing = await _recv_until(receiver_ws, "typing")
            assert typing["from_uid"] == sender_uid

    asyncio.run(run())


def test_ws_read_marks_and_notifies(app_server):
    sender_session, sender_name = _member()
    receiver_session, receiver_name = _member()
    sender_uid = _db_user(sender_name)["uid"]
    receiver_uid = _db_user(receiver_name)["uid"]

    sender_session.post(
        f"{BASE_URL}/messages/send",
        headers={"Accept": "application/json"},
        data={"content": "read me please", "receiver_uid": receiver_uid},
    )

    async def run():
        async with websockets.connect(
            WS_URL, additional_headers=_cookie_header(sender_session)
        ) as sender_ws, websockets.connect(
            WS_URL, additional_headers=_cookie_header(receiver_session)
        ) as receiver_ws:
            await _recv_until(sender_ws, "ready")
            await _recv_until(receiver_ws, "ready")

            await receiver_ws.send(
                json.dumps({"type": "read", "with_uid": sender_uid})
            )
            receipt = await _recv_until(sender_ws, "read")
            assert receipt["by_uid"] == receiver_uid

    asyncio.run(run())
    refresh_snapshot()
    rows = list(
        get_table("messages").find(
            sender_uid=sender_uid, receiver_uid=receiver_uid, read=True
        )
    )
    assert rows, "messages should be marked read after the read frame"