Spaces:
Running
Running
File size: 1,744 Bytes
f8381b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """Unit tests for the APIPoolManager rate limiter."""
from __future__ import annotations
import time
import pytest
from utils import APIPoolManager
def test_round_robin_no_limits() -> None:
pool = APIPoolManager(["k1", "k2", "k3"], rate_limits=None)
seen = [pool.get_next_key("any-model") for _ in range(6)]
# With no limits we should walk through all keys at least twice.
assert set(seen) == {"k1", "k2", "k3"}
def test_rpm_spacing_enforced() -> None:
"""With RPM=60 we expect a ~1s spacing between consecutive uses of the
same key. Two-key pool should let us avoid the wait."""
pool = APIPoolManager(["k1", "k2"], rate_limits={"m": (60, 1000)})
k_a = pool.get_next_key("m")
pool.record_usage(k_a, "m", time.time())
k_b = pool.get_next_key("m")
assert k_a != k_b, "Round-robin should pick the other key when one is hot"
def test_rpd_exhaustion_drops_key() -> None:
"""A key that hits its daily limit must be removed from active pool."""
pool = APIPoolManager(["k1", "k2"], rate_limits={"m": (60, 2)})
for _ in range(2):
k = pool.get_next_key("m")
pool.record_usage(k, "m")
# By now both keys may have hit their RPD=2. Next call should still work
# if at least one key has capacity, else raise RuntimeError.
keys_left = list(pool.active_keys)
if not keys_left:
with pytest.raises(RuntimeError):
pool.get_next_key("m")
else:
# Drain the remaining one too.
for _ in range(2):
try:
k = pool.get_next_key("m")
pool.record_usage(k, "m")
except RuntimeError:
break
assert not pool.active_keys, "Both keys should be exhausted now"
|