def test_different_user_ids_limited_separately(redis):
"""Ensure one user being rate-limited doesn't affect a different one."""
limit = 5
user_id = 1
action = RateLimitedAction("test", timedelta(hours=1), limit, redis=redis)
# check the action for the first user_id until it's blocked
result = action.check_for_user_id(user_id)
while result.is_allowed:
result = action.check_for_user_id(user_id)
# it should still be allowed for a different user_id
assert action.check_for_user_id(user_id + 1)
def test_simple_rate_limiting_by_user_id(redis):
"""Ensure simple rate-limiting by user_id is working."""
limit = 5
user_id = 1
# define an action with max_burst equal to the full limit
action = RateLimitedAction("testaction",
timedelta(hours=1),
limit,
max_burst=limit,
redis=redis)
# run the action the full number of times, should all be allowed
for _ in range(limit):
result = action.check_for_user_id(user_id)
assert result.is_allowed
# try one more time, should be rejected
result = action.check_for_user_id(user_id)
assert not result.is_allowed
def test_time_until_retry(redis):
"""Ensure an unbursted limit's time_until_retry is the expected value."""
user_id = 1
period = timedelta(seconds=60)
limit = 2
# create an action with no burst allowed, which will force the actions to be spaced
# "evenly" across the limit
action = RateLimitedAction("test",
period=period,
limit=limit,
max_burst=1,
redis=redis)
# first usage should be fine
result = action.check_for_user_id(user_id)
assert result.is_allowed
# second should fail, and require a wait of (period / limit) - 1 sec
result = action.check_for_user_id(user_id)
assert not result.is_allowed
assert result.time_until_retry == (period / limit) - timedelta(seconds=1)
def test_remaining_limit(redis):
"""Ensure a limit's "remaining limit" decreases as expected."""
user_id = 1
limit = 10
# create an action allowing the full limit as a burst
action = RateLimitedAction("test",
timedelta(days=1),
limit,
max_burst=limit,
redis=redis)
for count in range(1, limit + 1):
result = action.check_for_user_id(user_id)
assert result.remaining_limit == limit - count
def check_rate_limit(request: Request, action_name: str) -> RateLimitResult:
"""Check the rate limit for a particular action on a request."""
action = None
# check for a custom rate-limit for the user
if request.user:
user_limit = (
request.query(UserRateLimit)
.filter(
UserRateLimit.user == request.user, UserRateLimit.action == action_name
)
.one_or_none()
)
if user_limit:
action = RateLimitedAction(
action_name,
user_limit.period,
user_limit.limit,
by_user=True,
by_ip=False,
)
# if a custom rate-limit wasn't found, use the default, global rate-limit
if not action:
try:
action = RATE_LIMITED_ACTIONS[action_name]
except KeyError as exc:
raise ValueError("Invalid action name: %s" % action_name) from exc
action.redis = request.redis
results = []
if action.is_global:
results.append(action.check_global())
if action.by_user and request.user:
results.append(action.check_for_user_id(request.user.user_id))
if action.by_ip and request.remote_addr:
results.append(action.check_for_ip(request.remote_addr))
# no checks were done, return the "not limited" result
if not results:
return RateLimitResult.unlimited_result()
return RateLimitResult.merged_result(results)
def test_max_burst_defaults_to_half(redis):
"""Ensure that unspecified max_burst on a RateLimitedAction allows half."""
limit = 10
user_id = 1
action = RateLimitedAction("test", timedelta(days=1), limit, redis=redis)
# see how many times we can do the action until it gets blocked
count = 0
while True:
result = action.check_for_user_id(user_id)
if result.is_allowed:
count += 1
else:
break
assert count == limit // 2
def test_check_by_user_id_disabled():
"""Ensure non-by_user RateLimitedAction can't be checked by user_id."""
action = RateLimitedAction("test", timedelta(hours=1), 5, by_user=False)
with raises(RateLimitError):
action.check_for_user_id(1)