refactor: RateLimiter принимает _time_func для детерминированных тестов без зависимости от реального времени
This commit is contained in:
parent
3a2c388adf
commit
5eb4184d63
@ -1,45 +1,80 @@
|
||||
"""Тесты для utils/rate_limiter.py — проверка логики токен-бакета."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
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 = RateLimiter(rate=2.0, burst=5)
|
||||
limiter, _ = _make_time()
|
||||
assert limiter.tokens == 5.0
|
||||
|
||||
|
||||
async def test_acquire_consumes_token() -> None:
|
||||
"""acquire() уменьшает количество токенов."""
|
||||
limiter = RateLimiter(rate=1.0, burst=3)
|
||||
limiter, _ = _make_time()
|
||||
await limiter.acquire()
|
||||
assert limiter.tokens == 2.0
|
||||
assert limiter.tokens == 4.0
|
||||
|
||||
|
||||
async def test_acquire_waits_when_empty() -> None:
|
||||
"""acquire() ждёт, когда токены закончились."""
|
||||
limiter = RateLimiter(rate=10.0, burst=1) # 10 токенов/сек
|
||||
await limiter.acquire() # бакет пуст
|
||||
start = time.monotonic()
|
||||
await limiter.acquire() # должен ждать ~0.1 сек
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed >= 0.05 # допускаем погрешность
|
||||
"""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 = RateLimiter(rate=100.0, burst=3)
|
||||
await asyncio.sleep(0.1) # теоретически +10 токенов, но cap = 3
|
||||
limiter, times = _make_time()
|
||||
times[0] = 10.0 # теоретически +100 токенов, но cap = 5
|
||||
async with limiter.lock:
|
||||
limiter._refill()
|
||||
assert limiter.tokens == 3.0
|
||||
assert limiter.tokens == 5.0
|
||||
|
||||
|
||||
async def test_multiple_acquire() -> None:
|
||||
"""Можно забрать несколько токенов за раз."""
|
||||
limiter = RateLimiter(rate=1.0, burst=10)
|
||||
await limiter.acquire(token=5)
|
||||
assert limiter.tokens == 5.0
|
||||
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
|
||||
|
||||
@ -10,7 +10,7 @@ import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Final
|
||||
from typing import Callable, Final
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -18,21 +18,28 @@ logger = logging.getLogger(__name__)
|
||||
class RateLimiter:
|
||||
"""Токен-бакет: заполняется со скоростью rate токенов/сек, максимум burst."""
|
||||
|
||||
def __init__(self, rate: float, burst: int) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
rate: float,
|
||||
burst: int,
|
||||
_time_func: Callable[[], float] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
rate: Скорость пополнения токенов (токенов в секунду).
|
||||
burst: Максимальный размер бакета.
|
||||
_time_func: Функция получения времени (для тестов). По умолчанию time.monotonic.
|
||||
"""
|
||||
self.rate: float = rate
|
||||
self.burst: int = burst
|
||||
self.tokens: float = float(burst)
|
||||
self.lock: asyncio.Lock = asyncio.Lock()
|
||||
self._last_refill: float = time.monotonic()
|
||||
self._time_func = _time_func or time.monotonic
|
||||
self._last_refill: float = self._time_func()
|
||||
|
||||
def _refill(self) -> None:
|
||||
"""Пополнить токены за прошедшее время."""
|
||||
now: float = time.monotonic()
|
||||
now: float = self._time_func()
|
||||
elapsed: float = now - self._last_refill
|
||||
self.tokens = min(self.burst, self.tokens + elapsed * self.rate)
|
||||
self._last_refill = now
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user