"""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