def __init__(self, api_key: str, http_client: HTTPClient = None) -> None:
try:
api_key = os.environ[api_key]
except KeyError:
pass
self._key = api_key
if http_client is None:
self._client = HTTPClient()
else:
self._client = http_client
self._rate_limiter = MultiRateLimiter(
FixedWindowRateLimiter(600, 3000), FixedWindowRateLimiter(10, 50))
self._cached_data = {}
def _construct_limiters(self, limits: List[List[int]]):
# Creates the necessary FixedWindowRateLimiters from the rates in the headers
assert len(self._limiters) == 0
# Create the rate limiters
for permits, window in limits:
self._limiters.append(
FixedWindowRateLimiter(window_seconds=window,
window_permits=permits))
def _get_rate_limiter(self, platform: Platform, endpoint: str):
try:
limiter = self._rate_limiters[(platform, endpoint)]
except KeyError:
# TODO: Move to settings and don't force dev limits.
# TODO This is a circular import and instead this functionality should somehow be moved into settings.
from cassiopeia.configuration import settings
limits = settings.rate_limits
limiter = MultiRateLimiter(*([
FixedWindowRateLimiter(window, permits)
for permits, window in limits
] + [
FixedWindowRateLimiter(
*reversed(_METHOD_RATE_LIMITS[endpoint]))
]))
self._rate_limiters[(platform, endpoint)] = limiter
return limiter
def test_window_decorator_simple():
limiter = FixedWindowRateLimiter(SECONDS, PERMITS)
@limiter.limit
def call():
return True
assert call()
def test_window_permit_count():
limiter = FixedWindowRateLimiter(SECONDS, MANY_PERMITS)
assert limiter.permits_issued == 0
@limiter.limit
def call():
pass
for i in range(LARGE_VALUE_COUNT // 2):
with limiter:
pass
assert limiter.permits_issued == i + 1
limiter.reset_permits_issued()
assert limiter.permits_issued == 0
for i in range(LARGE_VALUE_COUNT // 2):
call()
assert limiter.permits_issued == i + 1
def test_window_acquire_across_windows():
from time import monotonic, sleep
limiter = FixedWindowRateLimiter(SECONDS, 1)
# This should take up two fixed windows for the first task, and the second shouldn't execute until the third (after 2 x SECONDS).
with limiter:
first = monotonic()
sleep(SECONDS * 1.25)
with limiter:
second = monotonic()
assert (SECONDS - EPSILON) * 3 >= second - first >= (SECONDS - EPSILON) * 2
def test_window_acquire_timing():
from time import monotonic
limiter = FixedWindowRateLimiter(SECONDS, PERMITS)
times = []
for _ in range(VALUE_COUNT):
with limiter:
times.append(monotonic())
start_indexes = [i for i in range(VALUE_COUNT) if i % PERMITS == 0]
last = -SECONDS
for index in start_indexes:
assert times[index] - last >= SECONDS - EPSILON
last = times[index]
def test_window_decorator_across_windows():
from time import monotonic, sleep
limiter = FixedWindowRateLimiter(SECONDS, 1)
@limiter.limit
def call():
t = monotonic()
sleep(SECONDS * 1.25)
return t
# This should take up two fixed windows for the first task, and the second shouldn't execute until the third (after 2 x SECONDS).
first = call()
second = call()
assert (SECONDS - EPSILON) * 3 >= second - first >= (SECONDS - EPSILON) * 2
def test_window_decorator_timing():
from time import monotonic
limiter = FixedWindowRateLimiter(SECONDS, PERMITS)
@limiter.limit
def call():
return monotonic()
times = []
for _ in range(VALUE_COUNT):
times.append(call())
start_indexes = [i for i in range(VALUE_COUNT) if i % PERMITS == 0]
last = -SECONDS
for index in start_indexes:
assert times[index] - last >= SECONDS - EPSILON
last = times[index]
def test_window_acquire_simple():
limiter = FixedWindowRateLimiter(SECONDS, PERMITS)
x = False
with limiter:
x = True
assert x