def test_check_for_ip_invalid_address():
"""Ensure RateLimitedAction.check_for_ip can't take an invalid IP."""
ip = '123.456.789.123'
action = RateLimitedAction('testaction', timedelta(hours=1), 10)
with raises(ValueError):
action.check_for_ip(ip)
def test_reset_for_ip_invalid_address():
"""Ensure RateLimitedAction.reset_for_ip can't take an invalid IP."""
ip = "123.456.789.123"
action = RateLimitedAction("testaction", timedelta(hours=1), 10)
with raises(ValueError):
action.reset_for_ip(ip)
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_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 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_action_with_all_types_disabled():
"""Ensure RateLimitedAction can't have both by_user and by_ip disabled."""
with raises(ValueError):
RateLimitedAction("test",
timedelta(hours=1),
5,
by_user=False,
by_ip=False)
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 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:
raise ValueError("Invalid action name: %s" % action_name)
action.redis = request.redis
results = []
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_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_simple_global_rate_limiting(redis):
"""Ensure simple global rate-limiting is working."""
limit = 5
# define an action with max_burst equal to the full limit
action = RateLimitedAction(
"testaction",
timedelta(hours=1),
limit,
max_burst=limit,
by_user=False,
by_ip=False,
redis=redis,
)
# run the action the full number of times, should all be allowed
for _ in range(limit):
result = action.check_global()
assert result.is_allowed
# try one more time, should be rejected
result = action.check_global()
assert not result.is_allowed
def test_check_global_disabled():
"""Ensure global check is disabled if action is by_user or by_ip."""
action = RateLimitedAction("test",
timedelta(hours=1),
5,
by_user=True,
by_ip=False)
with raises(RateLimitError):
action.check_global()
action = RateLimitedAction("test",
timedelta(hours=1),
5,
by_user=False,
by_ip=True)
with raises(RateLimitError):
action.check_global()
def test_max_burst_with_limit_1():
"""Ensure an action with limit 1 also has its max_burst set to 1."""
action = RateLimitedAction("test", timedelta(hours=1), 1)
assert action.max_burst == 1
def test_check_by_ip_disabled():
"""Ensure non-by_ip RateLimitedAction can't be checked by ip."""
action = RateLimitedAction("test", timedelta(hours=1), 5, by_ip=False)
with raises(RateLimitError):
action.check_for_ip("123.123.123.123")
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)