devplacepy/tests/api/messages/wsticket.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"
JSON = {"Accept": "application/json"}
_counter = [0]


@pytest.fixture(scope="module", autouse=True)
def _wsticket_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 _unique(prefix="wst"):
    _counter[0] += 1
    return f"{prefix}{int(time.time() * 1000)}{_counter[0]}"


def _signup():
    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",
        },
        allow_redirects=True,
    )
    return s, name


def _uid(name):
    refresh_snapshot()
    return get_table("users").find_one(username=name)["uid"]


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_ticket_requires_auth_401_for_json(app_server):
    r = requests.post(
        f"{BASE_URL}/messages/ws-ticket", headers=JSON, allow_redirects=False
    )
    assert r.status_code == 401


def test_ws_ticket_requires_auth_redirects_browser(app_server):
    r = requests.post(f"{BASE_URL}/messages/ws-ticket", allow_redirects=False)
    assert r.status_code == 303


def test_ws_ticket_issued_for_logged_in_session(app_server):
    session, _ = _signup()
    r = session.post(f"{BASE_URL}/messages/ws-ticket", headers=JSON)
    assert r.status_code == 200, r.text[:300]
    data = r.json()
    assert data["ticket"]
    assert data["expires_in"] == 30


def test_ws_ticket_single_use_and_authenticates_owner(app_server):
    session, name = _signup()
    owner_uid = _uid(name)

    ticket = session.post(f"{BASE_URL}/messages/ws-ticket", headers=JSON).json()[
        "ticket"
    ]

    async def first_connect():
        async with websockets.connect(f"{WS_URL}?ticket={ticket}") as ws:
            ready = await _recv_until(ws, "ready")
            assert ready["user_uid"] == owner_uid

    asyncio.run(first_connect())

    async def second_connect():
        async with websockets.connect(f"{WS_URL}?ticket={ticket}") as ws:
            with pytest.raises(websockets.exceptions.ConnectionClosed):
                await asyncio.wait_for(ws.recv(), timeout=3.0)

    asyncio.run(second_connect())