feedBack/tests/test_ws_sync.py

"""Tests for the session-sync relay WebSocket (/ws/sync/{session_id}).

Behavior tests run against a minimal FastAPI app carrying just the router
(fast — no full-server import); one integration test imports the real server
to pin that the route is actually mounted there.

Covers the feedBack#1030 acceptance list: bidirectional fan-out, late join,
sender never echoed, room garbage collection, and the limit closes (invalid
session id, binary frames, frame size, room size, room count, rate cap) —
including that one client tripping a limit doesn't disturb the others.
"""

from __future__ import annotations

import asyncio
import importlib
import sys
import time

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect

from routers import ws_sync


@pytest.fixture(autouse=True)
def _clean_rooms():
    ws_sync._rooms.clear()
    yield
    ws_sync._rooms.clear()


@pytest.fixture()
def client():
    app = FastAPI()
    app.include_router(ws_sync.router)
    with TestClient(app) as c:
        yield c


def _expect_close(ws, code):
    with pytest.raises(WebSocketDisconnect) as exc:
        ws.receive_text()
    assert exc.value.code == code


# ── Fan-out semantics ────────────────────────────────────────────────────────

def test_two_clients_relay_both_directions_and_no_echo(client):
    with client.websocket_connect("/ws/sync/ROOM01") as a, \
         client.websocket_connect("/ws/sync/ROOM01") as b:
        a.send_text('{"type":"time","t":1.5}')
        assert b.receive_text() == '{"type":"time","t":1.5}'
        b.send_text('{"type":"hello"}')
        # A's first inbound frame is B's hello — NOT an echo of its own send.
        assert a.receive_text() == '{"type":"hello"}'


def test_late_joiner_receives_subsequent_frames(client):
    with client.websocket_connect("/ws/sync/ROOM02") as a, \
         client.websocket_connect("/ws/sync/ROOM02") as b:
        a.send_text("f1")
        assert b.receive_text() == "f1"
        with client.websocket_connect("/ws/sync/ROOM02") as c:
            a.send_text("f2")
            assert b.receive_text() == "f2"
            assert c.receive_text() == "f2"


def test_rooms_are_isolated(client):
    with client.websocket_connect("/ws/sync/ROOMA1") as a, \
         client.websocket_connect("/ws/sync/ROOMB1") as b, \
         client.websocket_connect("/ws/sync/ROOMA1") as a2:
        a.send_text("for-room-a")
        assert a2.receive_text() == "for-room-a"
        # B (other room) got nothing: prove it by relaying within B's room.
        with client.websocket_connect("/ws/sync/ROOMB1") as b2:
            b2.send_text("for-room-b")
            assert b.receive_text() == "for-room-b"


def test_client_disconnect_does_not_disrupt_remaining(client):
    with client.websocket_connect("/ws/sync/ROOM03") as a, \
         client.websocket_connect("/ws/sync/ROOM03") as b:
        with client.websocket_connect("/ws/sync/ROOM03") as c:
            a.send_text("before")
            assert b.receive_text() == "before"
            assert c.receive_text() == "before"
        # C is gone; relay between A and B continues.
        a.send_text("after")
        assert b.receive_text() == "after"


def test_room_garbage_collected_when_last_client_leaves(client):
    with client.websocket_connect("/ws/sync/ROOM04") as a:
        with client.websocket_connect("/ws/sync/ROOM04") as b:
            a.send_text("x")
            assert b.receive_text() == "x"
            assert "ROOM04" in ws_sync._rooms
    assert "ROOM04" not in ws_sync._rooms
    assert ws_sync._rooms == {}


# ── Limit enforcement ────────────────────────────────────────────────────────

@pytest.mark.parametrize("bad_id", ["abc", "x" * 65, "has space", "bad$id", "nope!"])
def test_invalid_session_id_closed_with_policy_code(client, bad_id):
    with client.websocket_connect(f"/ws/sync/{bad_id}") as ws:
        _expect_close(ws, 1008)
    assert ws_sync._rooms == {}


def test_binary_frame_closes_with_unsupported_data(client):
    with client.websocket_connect("/ws/sync/ROOM05") as ws:
        ws.send_bytes(b"\x00\x01")
        _expect_close(ws, 1003)


def test_oversized_frame_closes_sender_only(client):
    with client.websocket_connect("/ws/sync/ROOM06") as a, \
         client.websocket_connect("/ws/sync/ROOM06") as b, \
         client.websocket_connect("/ws/sync/ROOM06") as c:
        a.send_text("x" * (ws_sync.MAX_FRAME_BYTES + 1))
        _expect_close(a, 1009)
        # The room carries on without A.
        b.send_text("still-alive")
        assert c.receive_text() == "still-alive"


def test_room_client_cap(client, monkeypatch):
    monkeypatch.setattr(ws_sync, "MAX_CLIENTS_PER_ROOM", 2)
    with client.websocket_connect("/ws/sync/ROOM07") as a, \
         client.websocket_connect("/ws/sync/ROOM07") as b, \
         client.websocket_connect("/ws/sync/ROOM07") as c:
        _expect_close(c, 1013)
        a.send_text("two-is-fine")
        assert b.receive_text() == "two-is-fine"


def test_total_room_cap(client, monkeypatch):
    monkeypatch.setattr(ws_sync, "MAX_ROOMS", 1)
    with client.websocket_connect("/ws/sync/ROOM08"):
        with client.websocket_connect("/ws/sync/ROOM09") as overflow:
            _expect_close(overflow, 1013)
        # Joining the EXISTING room is still fine at the room cap.
        with client.websocket_connect("/ws/sync/ROOM08"):
            pass


def test_rate_cap_closes_flooding_sender(client, monkeypatch):
    monkeypatch.setattr(ws_sync, "RATE_BURST", 3.0)
    monkeypatch.setattr(ws_sync, "RATE_MSGS_PER_SEC", 0.0)
    with client.websocket_connect("/ws/sync/ROOM10") as a, \
         client.websocket_connect("/ws/sync/ROOM10") as b:
        for i in range(3):
            a.send_text(f"burst-{i}")
        for i in range(3):
            assert b.receive_text() == f"burst-{i}"
        a.send_text("one-too-many")
        _expect_close(a, 1008)
        # The over-limit frame was dropped, not relayed, and B lives on.
        with client.websocket_connect("/ws/sync/ROOM10") as c:
            c.send_text("fresh-socket")
            assert b.receive_text() == "fresh-socket"


class _StalledPeer:
    """A fake room member whose send never completes (peer stopped draining)."""

    async def send_text(self, text):
        await asyncio.Event().wait()


def test_stalled_peer_is_evicted_and_healthy_peers_still_receive(client, monkeypatch):
    monkeypatch.setattr(ws_sync, "SEND_TIMEOUT_SECONDS", 0.2)
    with client.websocket_connect("/ws/sync/ROOM11") as a, \
         client.websocket_connect("/ws/sync/ROOM11") as b:
        # Wait for both handlers to have registered in the room, then inject
        # the stalled peer directly (a real stalled TCP peer isn't
        # constructible under TestClient).
        deadline = time.monotonic() + 2.0
        while len(ws_sync._rooms.get("ROOM11", {})) < 2:
            assert time.monotonic() < deadline, "room never filled"
            time.sleep(0.01)
        stalled = _StalledPeer()
        ws_sync._rooms["ROOM11"][stalled] = asyncio.Lock()

        # Healthy delivery is not blocked behind the stalled peer, and by the
        # time a second frame has round-tripped, the first fan-out's timeout
        # has fired and evicted it.
        a.send_text("f1")
        assert b.receive_text() == "f1"
        a.send_text("f2")
        assert b.receive_text() == "f2"
        assert stalled not in ws_sync._rooms["ROOM11"]


def test_main_run_caps_uvicorn_ws_max_size():
    """main.py must bound inbound WS frames at the transport (uvicorn defaults
    to 16 MB, which would let a client materialize frames far past the relay's
    16 KB application cap before the handler ever sees them)."""
    import unittest.mock

    import main

    with (
        unittest.mock.patch("logging_setup.configure_logging"),
        unittest.mock.patch("uvicorn.run") as mock_run,
    ):
        main.run()

    kwargs = mock_run.call_args.kwargs
    assert kwargs.get("ws_max_size") == 64 * 1024
    assert kwargs["ws_max_size"] >= ws_sync.MAX_FRAME_BYTES


# ── Real-app integration ─────────────────────────────────────────────────────

def test_route_mounted_on_real_server(tmp_path, monkeypatch):
    monkeypatch.setenv("CONFIG_DIR", str(tmp_path / "config"))
    monkeypatch.setenv("DLC_DIR", str(tmp_path / "dlc"))
    monkeypatch.setenv("FEEDBACK_SYNC_STARTUP", "1")
    sys.modules.pop("server", None)
    server = importlib.import_module("server")
    monkeypatch.setattr(server, "load_plugins", lambda *a, **kw: None)
    monkeypatch.setattr(server, "startup_scan", lambda: None)

    with TestClient(server.app) as client:
        with client.websocket_connect("/ws/sync/REALAPP") as a, \
             client.websocket_connect("/ws/sync/REALAPP") as b:
            a.send_text('{"type":"time","t":0}')
            assert b.receive_text() == '{"type":"time","t":0}'