"""Agent bandwidth governor (TokenBucket + PidReadMeter) — pure stdlib, so it runs in the unit lane even though the rest of the agent (torch/ultralytics) can't be imported in CI. Timing asserts use generous margins: they check the throttle's ORDER of magnitude, not scheduler precision.""" import os import threading import time from agent.fc_agent.throttle import PidReadMeter, TokenBucket def test_unlimited_never_blocks_and_counts(): b = TokenBucket(0) t0 = time.monotonic() for _ in range(100): b.take(10_000_000) assert time.monotonic() - t0 < 0.2 assert b.consumed == 1_000_000_000 def test_take_enforces_average_rate(): # 400 KB/s with a 1s burst: the first 400KB is free (burst), the next # 200KB is 0.5s of debt the caller must wait out. b = TokenBucket(400_000) b.take(400_000) t0 = time.monotonic() b.take(200_000) elapsed = time.monotonic() - t0 assert 0.35 <= elapsed < 3.0 def test_set_rate_zero_unblocks_a_waiter(): b = TokenBucket(1_000) # 1 KB/s → a 1MB take would wait ~17 min started = threading.Event() def blocked_take(): started.set() b.take(1_000_000) th = threading.Thread(target=blocked_take, daemon=True) th.start() started.wait(1.0) time.sleep(0.1) # let it enter the wait loop b.set_rate(0) # lift the cap live (the UI dial) th.join(timeout=2.0) assert not th.is_alive() def test_charge_reports_debt_without_blocking(): b = TokenBucket(1_000) t0 = time.monotonic() assert b.charge(1_000) == 0.0 # covered by the 1s burst debt = b.charge(2_000) # 2s over budget → ~2s of debt assert time.monotonic() - t0 < 0.2 assert 1.5 <= debt <= 2.1 assert b.consumed == 3_000 def test_pid_read_meter_tracks_own_reads(): m = PidReadMeter(os.getpid()) first = m.delta() # cumulative rchar since process start assert first is not None and first > 0 with open("/dev/zero", "rb") as f: f.read(1_048_576) grown = m.delta() assert grown is not None and grown >= 1_048_576 def test_pid_read_meter_missing_pid_degrades_to_none(): # PID 2^22+ is above the default pid_max — guaranteed absent. assert PidReadMeter(2**22 + 12345).delta() is None