Wizard/wizard_server/tests/test_limits.py

"""Unit tests of the limiter, load-monitor and AI-bridge protection building blocks."""

import asyncio
import random

from wizard_game import WizardGame

from app.ai_bridge import AIBridge, AIConnection
from app.config import Settings
from app.limits import ConnectionLimiter, KeyedRateLimiter, TokenBucket
from app.load import (
    LoadLevel,
    LoadMonitor,
    LoadSample,
    LoadThresholds,
    classify,
    next_level,
    own_cgroup_dir,
    parse_cpu_max,
    parse_cpu_stat,
)

N, B, O = LoadLevel.NORMAL, LoadLevel.BUSY, LoadLevel.OVERLOADED


class FakeClock:
    def __init__(self):
        self.now = 0.0

    def __call__(self):
        return self.now


# --- TokenBucket / KeyedRateLimiter / ConnectionLimiter -----------------------


def test_token_bucket_burst_and_refill():
    clock = FakeClock()
    bucket = TokenBucket(rate=2, burst=3, clock=clock)
    assert [bucket.allow() for _ in range(4)] == [True, True, True, False]
    clock.now = 0.5  # one token back
    assert bucket.allow() and not bucket.allow()
    clock.now = 100  # never more than the burst
    assert sum(bucket.allow() for _ in range(10)) == 3


def test_keyed_rate_limiter_is_per_key_and_prunes():
    clock = FakeClock()
    limiter = KeyedRateLimiter(rate=1, burst=1, clock=clock)
    assert limiter.allow("a") and not limiter.allow("a")
    assert limiter.allow("b")
    clock.now = 10
    limiter.prune()
    assert len(limiter) == 0


def test_connection_limiter():
    limiter = ConnectionLimiter(max_total=3, max_per_ip=2)
    assert limiter.acquire("1.1.1.1") and limiter.acquire("1.1.1.1")
    assert not limiter.acquire("1.1.1.1")
    assert limiter.acquire("2.2.2.2")
    assert not limiter.acquire("3.3.3.3")  # total reached
    limiter.release("1.1.1.1")
    assert limiter.total == 2
    assert limiter.acquire("3.3.3.3")
    limiter.release("9.9.9.9")  # unknown ip is ignored
    assert limiter.total == 3


# --- Load levels --------------------------------------------------------------

T = LoadThresholds(
    lag_busy=0.1, lag_overload=0.25, lag_severe=1.0, cpu_busy=0.7, cpu_overload=0.9,
    throttle_overload=0.5, overload_samples=3, recover_samples=4,
)


def test_classify():
    assert classify(LoadSample(0.0, 0.1, 0.0), T) is N
    assert classify(LoadSample(0.15), T) is B
    assert classify(LoadSample(0.0, 0.75), T) is B
    assert classify(LoadSample(0.3), T) is O
    assert classify(LoadSample(0.0, 0.95), T) is O
    assert classify(LoadSample(0.0, 0.2, 0.6), T) is O


def run_levels(samples):
    level, recent, levels = N, [], []
    for sample in samples:
        recent.append(sample)
        level = next_level(level, recent, T)
        levels.append(level)
    return levels


def test_overload_needs_consecutive_samples():
    assert run_levels([O, O, N]) == [B, B, B]
    assert run_levels([O, O, O]) == [B, B, O]


def test_overload_recovery_needs_a_calm_window():
    levels = run_levels([O, O, O, N, N, N, B, N, N, N, N])
    assert levels[3:6] == [O, O, O]  # still overloaded while recovering
    assert levels[6] is B  # 4 samples without overload -> busy (last sample busy)
    assert levels[-1] is N  # 4 normal samples -> normal


def test_severe_lag_overloads_immediately():
    assert next_level(N, [O], T, severe=True) is O
    monitor = LoadMonitor(Settings())
    monitor.record(LoadSample(loop_lag=2.0))
    assert monitor.level is O


def test_busy_is_immediate_and_sticky():
    assert run_levels([N, B, N, N, N, N]) == [N, B, B, B, B, N]


# --- cgroup parsing / LoadMonitor ---------------------------------------------


def test_parse_cgroup_files():
    assert parse_cpu_max("max 100000\n") is None
    assert parse_cpu_max("50000 100000\n") == 0.5
    assert parse_cpu_max("") is None
    stats = parse_cpu_stat("usage_usec 10\nnr_periods 200\nnr_throttled 50\n")
    assert stats["nr_periods"] == 200 and stats["nr_throttled"] == 50


def test_own_cgroup_dir(tmp_path):
    (tmp_path / "user.slice" / "app.scope").mkdir(parents=True)
    proc = tmp_path / "cgroup"
    proc.write_text("0::/user.slice/app.scope\n")
    assert own_cgroup_dir(tmp_path, proc) == tmp_path / "user.slice" / "app.scope"
    proc.write_text("0::/\n")  # inside a container
    assert own_cgroup_dir(tmp_path, proc) == tmp_path
    assert own_cgroup_dir(tmp_path, tmp_path / "missing") == tmp_path


def test_monitor_reads_cgroup_throttling(tmp_path):
    (tmp_path / "cpu.max").write_text("50000 100000\n")
    (tmp_path / "cpu.stat").write_text("nr_periods 100\nnr_throttled 0\n")
    monitor = LoadMonitor(Settings(), cgroup_dir=tmp_path)
    (tmp_path / "cpu.stat").write_text("nr_periods 200\nnr_throttled 80\n")
    sample = monitor._measure(lag=0.0)
    assert sample.throttled == 0.8


def test_monitor_levels_and_callback():
    changes = []
    monitor = LoadMonitor(Settings(load_overload_samples=2, load_recover_samples=2), on_change=changes.append)
    for _ in range(2):
        monitor.record(LoadSample(loop_lag=0.5))  # overloaded, but not severe
    assert monitor.level is O
    for _ in range(2):
        monitor.record(LoadSample(loop_lag=0.0))
    assert monitor.level is N
    assert changes == [B, O, N]
    assert monitor.stats()["load"] == "normal"


def test_forced_level():
    monitor = LoadMonitor(Settings(load_force_level="overloaded"))
    assert monitor.level is O
    monitor.record(LoadSample(loop_lag=0.0))
    assert monitor.level is O


# --- AI bridge protection -----------------------------------------------------


class SilentWebSocket:
    """Accepts requests but never answers (a hanging AI service)."""

    def __init__(self):
        self.sent = []

    async def send_json(self, message):
        self.sent.append(message)


def observation():
    game = WizardGame(3, seed=1)
    return game.observation(game.current_player)


def test_ai_circuit_breaker():
    async def scenario():
        bridge = AIBridge(timeout=0.01, rng=random.Random(0), breaker_failures=2, breaker_cooldown=60)
        ws = SilentWebSocket()
        bridge._connections.append(AIConnection(ws, "fake"))
        obs = observation()
        for _ in range(4):
            action = await bridge.decide(obs)
            assert action in obs.legal_actions  # random fallback
        assert len(ws.sent) == 2  # after 2 timeouts the AI is not asked anymore
        assert ws.sent[0]["timeout"] == 0.01
        assert bridge.stats()["paused"] is True

    asyncio.run(scenario())


class OverloadedAIWebSocket:
    """Answers every request immediately with the given error message."""

    def __init__(self, bridge_connection_ref, message):
        self.ref, self.message, self.sent = bridge_connection_ref, message, []

    async def send_json(self, request):
        self.sent.append(request)
        future = self.ref[0].pending[request["request_id"]]
        future.set_result({"type": "error", "request_id": request["request_id"], "message": self.message})


def test_load_shedding_does_not_trip_the_breaker():
    async def scenario():
        bridge = AIBridge(timeout=1, rng=random.Random(0), breaker_failures=2)
        ref = []
        ws = OverloadedAIWebSocket(ref, "overloaded")
        connection = AIConnection(ws, "fake")
        ref.append(connection)
        bridge._connections.append(connection)
        obs = observation()
        for _ in range(5):
            assert await bridge.decide(obs) in obs.legal_actions
        assert len(ws.sent) == 5  # still asked every time
        assert bridge.stats()["paused"] is False

        ws.message = "model crashed"  # a real failure counts
        for _ in range(3):
            await bridge.decide(obs)
        assert len(ws.sent) == 7 and bridge.stats()["paused"] is True

    asyncio.run(scenario())


def test_ai_inflight_limit():
    async def scenario():
        bridge = AIBridge(timeout=1, rng=random.Random(0), max_inflight=1)
        ws = SilentWebSocket()
        connection = AIConnection(ws, "fake")
        connection.pending["busy"] = asyncio.get_running_loop().create_future()
        bridge._connections.append(connection)
        obs = observation()
        for _ in range(10):
            assert await bridge.decide(obs) in obs.legal_actions
        assert ws.sent == []  # not even sent: too many requests already waiting
        assert bridge.stats()["paused"] is False  # load shedding is not a failure

    asyncio.run(scenario())