devplacepy/tests/unit/services/nodes/session.py

# retoor <retoor@molodetz.nl>

import asyncio
import json

import pytest

from devplacepy.services.nodes.frames import FrameError
from devplacepy.services.nodes.registry import NodeCounters, NodeRegistry
from devplacepy.services.nodes.session import END_HEARTBEAT, NodeSession
from devplacepy.services.nodes.spec import CLOSE_KICKED, CLOSE_POLICY, SPEC
from tests.conftest import run_async as run

LIMITS = SPEC.limits
HELLO = {
    "type": "hello",
    "protocol": SPEC.version,
    "client_version": "0.1.0",
    "kind": "native",
    "platform": "linux/amd64",
    "modes": ["tunnel", "fetch"],
}


class FakeSocket:
    def __init__(self, incoming=()):
        self.incoming: asyncio.Queue = asyncio.Queue()
        for item in incoming:
            self.incoming.put_nowait(item)
        self.sent: list[dict] = []
        self.closed: tuple[int, str] | None = None

    async def receive_text(self):
        return await self.incoming.get()

    async def send_text(self, text):
        self.sent.append(json.loads(text))

    async def close(self, code=1000, reason=""):
        self.closed = (code, reason)


def session_for(socket, user="u1", uid="d1", kind="native"):
    return NodeSession(socket, {"uid": uid, "user_uid": user, "kind": kind})




def test_handshake_sends_welcome(local_db):
    socket = FakeSocket([json.dumps(HELLO)])
    session = session_for(socket)
    frame = run(session.handshake())
    assert frame["platform"] == "linux/amd64"
    welcome = socket.sent[0]
    assert welcome["type"] == "welcome" and welcome["device_uid"] == "d1"
    assert welcome["limits"] == LIMITS.describe()


@pytest.mark.parametrize(
    "mutation,code",
    [
        ({"protocol": SPEC.version + 1}, "protocol_unsupported"),
        ({"protocol": SPEC.min_supported - 1}, "protocol_unsupported"),
        ({"kind": "extension"}, "kind_mismatch"),
        ({"modes": []}, "bad_modes"),
        ({"modes": ["tunnel", "teleport"]}, "bad_modes"),
    ],
)
def test_handshake_rejects_bad_hellos(local_db, mutation, code):
    socket = FakeSocket([json.dumps({**HELLO, **mutation})])
    with pytest.raises(FrameError) as info:
        run(session_for(socket).handshake())
    assert info.value.code == code
    assert socket.sent == []


def test_the_first_frame_must_be_hello(local_db):
    socket = FakeSocket([json.dumps({"type": "pong"})])
    with pytest.raises(FrameError) as info:
        run(session_for(socket).handshake())
    assert info.value.code == "expected_hello"


def test_a_silent_client_times_out(local_db, monkeypatch):
    monkeypatch.setattr("devplacepy.services.nodes.session.LIMITS", type("L", (), {
        "hello_timeout_seconds": 0.05, "ping_seconds": 1, "missed_pings_before_evict": 2,
        "describe": staticmethod(lambda: {}),
    }))
    with pytest.raises(FrameError) as info:
        run(session_for(FakeSocket()).handshake())
    assert info.value.code == "hello_timeout" and info.value.close_code == CLOSE_POLICY


def test_usage_is_advisory_and_summed():
    session = session_for(FakeSocket())
    run(session._handle({"type": "usage", "bytes_up": 5, "bytes_down": 7}))
    run(session._handle({"type": "usage", "bytes_up": -9, "bytes_down": 1}))
    assert (session.reported_up, session.reported_down) == (5, 8)


@pytest.mark.parametrize("kind", ["job_result", "job_error"])
def test_frames_without_an_assigned_job_are_protocol_errors(kind):
    with pytest.raises(FrameError):
        run(session_for(FakeSocket())._handle({"type": kind, "job_id": "nope"}))


@pytest.mark.parametrize("kind", ["tunnel_data", "tunnel_close", "tunnel_opened"])
def test_late_tunnel_frames_for_a_closed_relay_are_ignored(kind):
    run(session_for(FakeSocket())._handle({"type": kind, "conn_id": "gone"}))


def test_missed_pings_evict_after_the_floor(monkeypatch):
    monkeypatch.setattr("devplacepy.services.nodes.session.LIMITS", type("L", (), {
        "ping_seconds": 0.02, "missed_pings_before_evict": 2,
    }))
    socket = FakeSocket()
    session = session_for(socket)
    ended = run(session.run())
    assert ended == END_HEARTBEAT
    assert [item["type"] for item in socket.sent] == ["ping", "ping"]
    assert socket.closed[0] == CLOSE_POLICY


def test_any_frame_resets_the_missed_counter(monkeypatch):
    monkeypatch.setattr("devplacepy.services.nodes.session.LIMITS", type("L", (), {
        "ping_seconds": 0.05, "missed_pings_before_evict": 2,
    }))
    socket = FakeSocket()
    session = session_for(socket)

    async def scenario():
        task = asyncio.create_task(session.run())
        await asyncio.sleep(0.07)
        assert session.missed_pings == 1
        await socket.incoming.put(json.dumps({"type": "pong"}))
        await asyncio.sleep(0.02)
        assert session.missed_pings == 0
        task.cancel()

    run(scenario())


def test_kick_sends_the_reason_then_closes_once():
    socket = FakeSocket()
    session = session_for(socket)
    run(session.kick("revoked"))
    run(session.kick("again"))
    assert socket.sent == [{"type": "kick", "reason": "revoked"}]
    assert socket.closed == (CLOSE_KICKED, "revoked")


def test_registry_replaces_the_same_device_and_enforces_capacity(local_db):
    registry = NodeRegistry(NodeCounters())
    first = session_for(FakeSocket(), uid="d1")
    assert registry.admit(first) == (True, None)
    second = session_for(FakeSocket(), uid="d1")
    admitted, previous = registry.admit(second)
    assert admitted and previous is first
    registry.release(first)
    assert registry.get("d1") is second
    registry.release(second)
    assert registry.count() == 0
    counts = registry.counters.snapshot()
    assert counts["connected_total"] == 2 and counts["replaced"] == 1


def test_registry_capacity_zero_admits_nobody(local_db):
    from devplacepy.database import get_setting, set_setting

    previous = get_setting("nodes_max_total", "")
    set_setting("nodes_max_total", "0")
    try:
        registry = NodeRegistry(NodeCounters())
        assert registry.capacity() == 0
        assert registry.admit(session_for(FakeSocket())) == (False, None)
        assert registry.counters.snapshot()["rejected_capacity"] == 1
    finally:
        set_setting("nodes_max_total", previous)


def test_registry_capacity_never_exceeds_the_spec_ceiling(local_db):
    from devplacepy.database import get_setting, set_setting

    previous = get_setting("nodes_max_total", "")
    set_setting("nodes_max_total", str(LIMITS.max_connected_nodes * 10))
    try:
        assert NodeRegistry(NodeCounters()).capacity() == LIMITS.max_connected_nodes
    finally:
        set_setting("nodes_max_total", previous)


def test_counters_reject_unknown_names():
    with pytest.raises(KeyError):
        NodeCounters().inc("made_up")


def test_kick_all_empties_the_registry(local_db):
    registry = NodeRegistry(NodeCounters())
    sockets = [FakeSocket() for _ in range(3)]
    for index, socket in enumerate(sockets):
        registry.admit(session_for(socket, uid=f"d{index}"))
    assert run(registry.kick_all("bye")) == 3
    assert registry.count() == 0
    assert all(socket.closed[0] == CLOSE_KICKED for socket in sockets)