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