devplacepy/tests/api/nodes/ws.py

# retoor <retoor@molodetz.nl>

import json

import pytest
from starlette.websockets import WebSocketDisconnect

from devplacepy.services.manager import service_manager
from devplacepy.services.nodes.spec import (
    CLOSE_CAPACITY,
    CLOSE_KICKED,
    CLOSE_POLICY,
    CLOSE_PROTOCOL,
    CLOSE_RETRY_OWNER,
    SPEC,
)
from tests.api.nodes.devices import create
from tests.api.nodes.ticket import bearer


def hello(**overrides):
    frame = {
        "type": "hello",
        "protocol": SPEC.version,
        "client_version": "0.1.0",
        "kind": "native",
        "platform": "linux/amd64",
        "modes": ["tunnel"],
    }
    frame.update(overrides)
    return json.dumps(frame)


def ticket_for(client, member):
    body = create(client, member).json()
    ticket = client.post("/nodes/ticket", headers=bearer(body["secret"])).json()["ticket"]
    return body["device"]["uid"], body["secret"], ticket


def closed_with(connection):
    with pytest.raises(WebSocketDisconnect) as info:
        connection.receive_text()
    return info.value.code


def test_handshake_welcomes_with_the_published_numbers(client, member):
    device_uid, _, ticket = ticket_for(client, member)
    with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as ws:
        ws.send_text(hello())
        welcome = json.loads(ws.receive_text())
        assert welcome["type"] == "welcome"
        assert welcome["device_uid"] == device_uid
        assert welcome["limits"] == SPEC.limits.describe()
        assert {item["suffix"] for item in welcome["destinations"]} == {
            item.suffix for item in SPEC.destinations
        }
        assert welcome["settings"]["enabled"] is True
        service = service_manager.get_service("nodes")
        assert service.registry.get(device_uid) is not None
        ws.send_text(json.dumps({"type": "pong"}))
        ws.send_text(json.dumps({"type": "usage", "bytes_up": 5, "bytes_down": 7}))
    assert service.registry.get(device_uid) is None


def test_a_ticket_works_once(client, member):
    _, _, ticket = ticket_for(client, member)
    with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as ws:
        ws.send_text(hello())
        ws.receive_text()
    with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as again:
        assert closed_with(again) == CLOSE_POLICY


def test_no_ticket_and_a_chat_ticket_are_refused(client, member):
    with client.websocket_connect("/nodes/ws") as ws:
        assert closed_with(ws) == CLOSE_POLICY
    from devplacepy.services.messaging.tickets import issue_ticket

    chat_ticket = issue_ticket(member["uid"])
    with client.websocket_connect(f"/nodes/ws?ticket={chat_ticket}") as ws:
        assert closed_with(ws) == CLOSE_POLICY


def test_a_non_owner_worker_says_retry(client, member):
    _, _, ticket = ticket_for(client, member)
    service_manager.set_lock_owner(False)
    try:
        with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as ws:
            assert closed_with(ws) == CLOSE_RETRY_OWNER
    finally:
        service_manager.set_lock_owner(True)


def test_protocol_violations_close_loudly(client, member):
    cases = [
        ("not json", CLOSE_PROTOCOL),
        (json.dumps({"type": "pong"}), CLOSE_PROTOCOL),
        (hello(protocol=SPEC.version + 5), CLOSE_POLICY),
        (hello(kind="extension"), CLOSE_POLICY),
        (hello(modes=["bogus"]), CLOSE_PROTOCOL),
        (hello(extra="x"), CLOSE_PROTOCOL),
    ]
    for frame, expected in cases:
        _, _, ticket = ticket_for(client, make_fresh(member))
        with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as ws:
            ws.send_text(frame)
            assert closed_with(ws) == expected, frame


def make_fresh(member):
    from tests.api.nodes.conftest import make_member

    return make_member()


def test_an_unsolicited_job_result_is_a_protocol_error(client, member):
    _, _, ticket = ticket_for(client, member)
    with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as ws:
        ws.send_text(hello())
        ws.receive_text()
        ws.send_text(
            json.dumps(
                {"type": "job_result", "job_id": "j", "status": 200, "body_b64": ""}
            )
        )
        assert closed_with(ws) == CLOSE_PROTOCOL


def test_a_second_connection_of_the_same_device_replaces_the_first(client, member):
    device_uid, secret, first_ticket = ticket_for(client, member)
    second_ticket = client.post("/nodes/ticket", headers=bearer(secret)).json()["ticket"]
    with client.websocket_connect(f"/nodes/ws?ticket={first_ticket}") as first:
        first.send_text(hello())
        first.receive_text()
        with client.websocket_connect(f"/nodes/ws?ticket={second_ticket}") as second:
            second.send_text(hello())
            assert json.loads(second.receive_text())["type"] == "welcome"
            kicked = json.loads(first.receive_text())
            assert kicked == {"type": "kick", "reason": "replaced"}
            assert closed_with(first) == CLOSE_KICKED


def test_capacity_is_a_hard_ceiling(client, member):
    from devplacepy.database import get_setting, set_setting

    service = service_manager.get_service("nodes")
    _, _, ticket = ticket_for(client, member)
    previous = get_setting("nodes_max_total", "")
    set_setting("nodes_max_total", "0")
    try:
        with client.websocket_connect(f"/nodes/ws?ticket={ticket}") as ws:
            assert closed_with(ws) == CLOSE_CAPACITY
        assert service.counters.snapshot()["rejected_capacity"] >= 1
    finally:
        set_setting("nodes_max_total", previous)