undercooked/tests/unit/test_send_queue_policy.py

"""Unit: the send-queue policy and the per-connection guardrail arithmetic
(phase 06 task 02) — the `Connection`'s queue is the only buffer between
the game loop and a slow reader, so it must be bounded and honest:

* **snapshot droppability under a slow reader** — once the queue is at
  `MAX_SEND_QUEUE` a new frame may only enter after the *oldest snapshot*
  is dropped; `lobby` / `round` / `err` frames are never dropped and keep
  their order,
* **the one-`rate_limited`-err bound** — a flood queues at most one
  `rate_limited` `err` (the client is told once; the queue cannot grow
  with the flood),
* **no unbounded growth** — the queue caps at `MAX_SEND_QUEUE` for the
  dominant (snapshot) traffic; only the never-dropped frames can push it
  past that, and only by their own count (game-bounded: one
  `lobby` / `round` per state change, one coalesced `rate_limited`),
* **the token bucket** — `rate` tokens per second, a burst of `burst`,
  refilled strictly off the injected `Clock` (never a wall clock),
* **the abuse tracker** — 3 *consecutive* rate/oversize violations inside
  the 10 s window close `4404` (a clean frame resets the run); 5 shape
  violations inside the window close it (a violation older than the
  window does not count).
"""

from __future__ import annotations

import asyncio
import contextlib
import json
from typing import Any

import pytest

from undercooked.server.clock import FakeClock
from undercooked.server.protocol import (
    ErrMsg,
    ErrorCode,
    LobbyMsg,
    RoundMsg,
    SnapMsg,
)
from undercooked.server.ws import (
    ABUSE_WINDOW_S,
    MAX_SEND_QUEUE,
    AbuseTracker,
    Connection,
    TokenBucket,
    pump_out,
)

pytestmark = pytest.mark.unit


class _FakeWs:
    """A `WsLike` that records what the pump wrote (no socket)."""

    def __init__(self) -> None:
        self.sent: list[dict[str, Any]] = []

    async def send_text(self, data: str) -> None:
        self.sent.append(json.loads(data))


def _conn(clock: FakeClock) -> tuple[Connection, _FakeWs]:
    ws = _FakeWs()
    return (
        Connection(
            ws=ws,
            token="tok-unit",  # noqa: S106 - a test identity, not a credential
            player_id="tok-unit",
            seat_index=0,
            role="player",
            rate=TokenBucket(10.0, 10, clock),
        ),
        ws,
    )


def _snap(index: int) -> SnapMsg:
    return SnapMsg(type="snap", raw={"t": index, "p": []})


def _lobby(index: int) -> LobbyMsg:
    return LobbyMsg(
        type="lobby",
        code=f"C{index}",
        state="lobby",
        can_start=False,
        seats=[],
        spectators=[],
    )


def _round() -> RoundMsg:
    return RoundMsg(
        type="round",
        result="won",
        score=1,
        target=1,
        stars=0,
        time_left=0.0,
        counters={},
        longest_streak=0,
        players=[],
        level_id="x",
        seed=1,
        ticks=0,
    )


# ------------------------------------------------------------------ drop policy


async def test_slow_reader_drops_oldest_snapshots() -> None:
    clock = FakeClock()
    conn, _ws = _conn(clock)
    # A slow reader: nothing is pumped yet, so the queue fills.
    for index in range(3):
        conn.enqueue(_lobby(index))  # non-droppable
    for index in range(20):
        conn.enqueue(_snap(index))  # droppable
    assert conn.send_queue.qsize() <= MAX_SEND_QUEUE
    # Drain and check: the lobbies are all there, in order; the survivors
    # are the NEWEST snapshots.
    got: list[Any] = []
    while not conn.send_queue.empty():
        got.append(conn.send_queue.get_nowait())
    types = [item.type for item in got]
    assert types == ["lobby"] * 3 + ["snap"] * (MAX_SEND_QUEUE - 3)
    survivors = [item.raw["t"] for item in got if item.type == "snap"]
    assert survivors == list(
        range(20 - (MAX_SEND_QUEUE - 3), 20)
    ), "the oldest snapshots must be the ones dropped"


async def test_non_droppable_frames_preserve_order() -> None:
    clock = FakeClock()
    conn, _ws = _conn(clock)
    frames = [
        _lobby(0),
        _snap(0),
        _round(),
        _snap(1),
        _lobby(1),
        _snap(2),
        _round(),
        _snap(3),
        _lobby(2),
        _snap(4),
        _snap(5),
        _snap(6),
        _snap(7),
        _snap(8),
    ]
    for frame in frames:
        conn.enqueue(frame)
    got: list[Any] = []
    while not conn.send_queue.empty():
        got.append(conn.send_queue.get_nowait())
    # Nothing non-droppable was dropped; the snaps that survived are the
    # newest, and every frame keeps its relative order.
    non_droppable = [item for item in got if item.type in ("lobby", "round")]
    assert list(non_droppable) == [
        _lobby(0),
        _round(),
        _lobby(1),
        _round(),
        _lobby(2),
    ]
    snaps = [item for item in got if item.type == "snap"]
    assert [s.raw["t"] for s in snaps] == sorted(s.raw["t"] for s in snaps)


async def test_queue_cannot_grow_with_a_rate_limited_flood() -> None:
    clock = FakeClock()
    conn, _ws = _conn(clock)
    for index in range(50):  # a flood of refusals (e.g. 50 rate_limited errs)
        conn.enqueue(
            ErrMsg(
                type="err",
                code=ErrorCode.RATE_LIMITED,
                msg=f"refusal {index}",
            )
        )
    # At most ONE rate_limited err is queued (coalesced); the queue is
    # bounded regardless of the flood's size.
    held: list[Any] = []
    while True:
        try:
            held.append(conn.send_queue.get_nowait())
        except asyncio.QueueEmpty:
            break
    rate_limited = [item for item in held if item.code == ErrorCode.RATE_LIMITED]
    assert len(rate_limited) == 1
    assert len(held) <= MAX_SEND_QUEUE + 1


async def test_queue_caps_at_max_under_snapshot_flood() -> None:
    clock = FakeClock()
    conn, _ws = _conn(clock)
    for index in range(100):  # the dominant traffic, a slow reader: the
        conn.enqueue(_snap(index))  # queue holds the newest, caps at the max
    assert conn.send_queue.qsize() == MAX_SEND_QUEUE
    held: list[Any] = []
    while True:
        try:
            held.append(conn.send_queue.get_nowait())
        except asyncio.QueueEmpty:
            break
    assert [item.raw["t"] for item in held] == list(range(92, 100))


async def test_all_non_droppable_frames_are_delivered_never_dropped() -> None:
    clock = FakeClock()
    conn, _ws = _conn(clock)
    for index in range(30):
        conn.enqueue(_lobby(index))
    # Nothing is ever dropped: once the queue holds no snapshot there is
    # nothing to drop, and the queue grows with the non-droppable frames
    # the server actually sent (bounded by the game's own event count —
    # a lobby/round/err per state change, a coalesced `rate_limited`).
    assert conn.send_queue.qsize() == 30


async def test_pump_delivers_frames_in_order_to_a_slow_reader() -> None:
    clock = FakeClock()
    conn, ws = _conn(clock)
    for index in range(15):
        conn.enqueue(_snap(index))
    assert conn.send_queue.qsize() == MAX_SEND_QUEUE
    pump = asyncio.create_task(pump_out(ws, conn))
    await asyncio.sleep(0)  # let the pump run its passes
    conn.alive = False
    pump.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        await pump
    sent_snaps = [frame["raw"]["t"] for frame in ws.sent if frame["type"] == "snap"]
    # Whatever the reader saw is in order and a suffix of the batch (only
    # the oldest may have been dropped).
    assert sent_snaps == sorted(sent_snaps)
    assert sent_snaps == list(range(15 - len(sent_snaps), 15))


# ----------------------------------------------------------------- token bucket


def test_bucket_serves_its_burst_then_empty() -> None:
    clock = FakeClock()
    bucket = TokenBucket(rate=5.0, burst=5, clock=clock)
    assert all(bucket.take() for _ in range(5))
    assert not bucket.take()  # the sixth, with no time elapsed: empty


def test_bucket_refills_off_the_injected_clock() -> None:
    clock = FakeClock()
    bucket = TokenBucket(rate=5.0, burst=5, clock=clock)
    for _ in range(5):
        bucket.take()
    clock.advance(0.2)  # one token (5/s * 0.2 s)
    assert bucket.take()
    assert not bucket.take()
    clock.advance(1.0)  # a full second: back to the burst cap
    assert all(bucket.take() for _ in range(5))
    assert not bucket.take()


def test_bucket_never_refills_past_the_burst_cap() -> None:
    clock = FakeClock()
    bucket = TokenBucket(rate=5.0, burst=3, clock=clock)
    clock.advance(100.0)  # a very long idle
    assert bucket.tokens() == 3.0
    assert all(bucket.take() for _ in range(3))
    assert not bucket.take()


def test_bucket_rejects_impossible_parameters() -> None:
    clock = FakeClock()
    with pytest.raises(ValueError, match=r"rate must be > 0"):
        TokenBucket(rate=0.0, burst=1, clock=clock)
    with pytest.raises(ValueError, match=r"burst must be >= 1"):
        TokenBucket(rate=1.0, burst=0, clock=clock)


def test_bucket_tokens_reports_the_refilled_count() -> None:
    clock = FakeClock()
    bucket = TokenBucket(rate=2.0, burst=10, clock=clock)
    for _ in range(5):
        bucket.take()
    assert bucket.tokens() == 5.0
    clock.advance(0.5)  # +1 (2/s * 0.5 s)
    assert bucket.tokens() == 6.0


# ------------------------------------------------------------- the abuse tracker


def test_rate_close_after_three_consecutive_violations() -> None:
    clock = FakeClock()
    tracker = AbuseTracker(clock)
    assert not tracker.note_rate_violation()
    assert not tracker.note_rate_violation()
    assert tracker.note_rate_violation()  # the third, inside the window


def test_clean_frame_resets_the_consecutive_run() -> None:
    clock = FakeClock()
    tracker = AbuseTracker(clock)
    assert not tracker.note_rate_violation()
    tracker.note_clean()
    assert not tracker.note_rate_violation()
    assert not tracker.note_rate_violation()  # the run only has two now
    clock.advance(1.0)
    assert tracker.note_rate_violation()  # three consecutive again


def test_rate_violations_outside_the_window_do_not_close() -> None:
    clock = FakeClock()
    tracker = AbuseTracker(clock)
    assert not tracker.note_rate_violation()  # t=0
    clock.advance(ABUSE_WINDOW_S + 0.001)  # the first stamp falls out
    assert not tracker.note_rate_violation()  # t=10.001: the run is two, one in-window
    assert not tracker.note_rate_violation()  # t=10.001: the run is three, two in-window
    clock.advance(0.1)
    assert tracker.note_rate_violation()  # now three fit in the window


def test_shape_close_after_five_in_the_window() -> None:
    clock = FakeClock()
    tracker = AbuseTracker(clock)
    for _ in range(4):
        assert not tracker.note_shape_violation()
    assert tracker.note_shape_violation()  # the fifth


def test_shape_violations_outside_the_window_do_not_close() -> None:
    clock = FakeClock()
    tracker = AbuseTracker(clock)
    for _ in range(4):
        assert not tracker.note_shape_violation()  # four at t=0
    assert tracker.note_shape_violation()  # the fifth closes (all in-window)
    # Age the whole batch past the window, then it must not count again.
    clock.advance(ABUSE_WINDOW_S + 0.001)
    for _ in range(4):
        assert not tracker.note_shape_violation()  # one..four in-window
    assert tracker.note_shape_violation()  # the fifth in-window