59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from yuxi.channels.adapters.wechat.rate_limiter import TokenBucketRateLimiter
|
||
|
|
|
||
|
|
|
||
|
|
class TestTokenBucketRateLimiter:
|
||
|
|
def test_initial_tokens_available(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=10, per=60.0)
|
||
|
|
assert limiter.available_tokens == 10.0
|
||
|
|
|
||
|
|
def test_acquire_sync_depletes(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=5, per=60.0)
|
||
|
|
results = [limiter.acquire_sync() for _ in range(5)]
|
||
|
|
assert all(results)
|
||
|
|
assert limiter.acquire_sync() is False
|
||
|
|
|
||
|
|
def test_acquire_sync_beyond_limit(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=1, per=60.0)
|
||
|
|
assert limiter.acquire_sync() is True
|
||
|
|
assert limiter.acquire_sync() is False
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_acquire_async_depletes(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=3, per=60.0)
|
||
|
|
results = [await limiter.acquire() for _ in range(3)]
|
||
|
|
assert all(results)
|
||
|
|
assert await limiter.acquire() is False
|
||
|
|
|
||
|
|
def test_rate_property(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=20, per=30.0)
|
||
|
|
assert limiter.rate == 20
|
||
|
|
|
||
|
|
def test_period_property(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=10, per=45.0)
|
||
|
|
assert limiter.period == 45.0
|
||
|
|
|
||
|
|
def test_zero_rate(self):
|
||
|
|
limiter = TokenBucketRateLimiter(rate=0, per=60.0)
|
||
|
|
assert limiter.acquire_sync() is False
|
||
|
|
|
||
|
|
def test_tokens_refill_after_time(self):
|
||
|
|
import time
|
||
|
|
limiter = TokenBucketRateLimiter(rate=10, per=60.0)
|
||
|
|
for _ in range(10):
|
||
|
|
limiter.acquire_sync()
|
||
|
|
assert limiter.acquire_sync() is False
|
||
|
|
limiter._last_refill = time.monotonic() - 60.0
|
||
|
|
assert limiter.acquire_sync() is True
|
||
|
|
|
||
|
|
def test_tokens_never_exceed_rate(self):
|
||
|
|
import time
|
||
|
|
limiter = TokenBucketRateLimiter(rate=5, per=60.0)
|
||
|
|
limiter._last_refill = time.monotonic() - 600.0
|
||
|
|
limiter.acquire_sync()
|
||
|
|
assert limiter.available_tokens <= limiter.rate
|