81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
"""Тесты для utils/rate_limiter.py — проверка логики токен-бакета."""
|
||
|
||
import asyncio
|
||
from unittest.mock import MagicMock
|
||
|
||
from utils.rate_limiter import RateLimiter
|
||
|
||
|
||
def _make_time() -> tuple[RateLimiter, list[float]]:
|
||
"""Создать RateLimiter с контролируемой временной функцией."""
|
||
times: list[float] = [0.0]
|
||
|
||
def controlled_time() -> float:
|
||
return times[0]
|
||
|
||
limiter = RateLimiter(rate=10.0, burst=5, _time_func=controlled_time)
|
||
return limiter, times
|
||
|
||
|
||
async def test_initial_tokens_full() -> None:
|
||
"""Бакет заполнен до burst при создании."""
|
||
limiter, _ = _make_time()
|
||
assert limiter.tokens == 5.0
|
||
|
||
|
||
async def test_acquire_consumes_token() -> None:
|
||
"""acquire() уменьшает количество токенов."""
|
||
limiter, _ = _make_time()
|
||
await limiter.acquire()
|
||
assert limiter.tokens == 4.0
|
||
|
||
|
||
async def test_acquire_waits_when_empty() -> None:
|
||
"""acquire() ждёт, когда токены закончились (контролируемое время)."""
|
||
limiter, times = _make_time()
|
||
# Потратить все 5 токенов
|
||
for _ in range(5):
|
||
await limiter.acquire()
|
||
assert limiter.tokens < 1.0
|
||
|
||
# Пропустить 0.2 сек -> должно пополниться 2 токена (rate=10)
|
||
times[0] = 0.2
|
||
async with limiter.lock:
|
||
limiter._refill()
|
||
assert limiter.tokens >= 2.0
|
||
|
||
|
||
async def test_burst_cap() -> None:
|
||
"""Токены не превышают burst после долгого простоя."""
|
||
limiter, times = _make_time()
|
||
times[0] = 10.0 # теоретически +100 токенов, но cap = 5
|
||
async with limiter.lock:
|
||
limiter._refill()
|
||
assert limiter.tokens == 5.0
|
||
|
||
|
||
async def test_multiple_acquire() -> None:
|
||
"""Можно забрать несколько токенов за раз."""
|
||
limiter, _ = _make_time()
|
||
await limiter.acquire(token=3)
|
||
assert limiter.tokens == 2.0
|
||
|
||
|
||
async def test_refill_partial() -> None:
|
||
"""Пополнение за малый интервал времени."""
|
||
limiter, times = _make_time()
|
||
times[0] = 0.1 # 10 токенов/сек * 0.1 сек = 1 токен
|
||
async with limiter.lock:
|
||
limiter._refill()
|
||
assert limiter.tokens == 5.0 # был 5 + 1 = 6, но cap = 5
|
||
|
||
|
||
async def test_refill_exact() -> None:
|
||
"""Точное пополнение при частичном бакете."""
|
||
limiter, times = _make_time()
|
||
await limiter.acquire(token=3) # осталось 2
|
||
times[0] = 0.1 # +1 токен
|
||
async with limiter.lock:
|
||
limiter._refill()
|
||
assert limiter.tokens == 3.0 # 2 + 1 = 3
|