def test_merged_results_single():
"""Ensure "merging" a single result just returns the same one."""
result = RateLimitResult(
is_allowed=True,
total_limit=50,
remaining_limit=22,
time_until_max=timedelta(seconds=256),
)
assert RateLimitResult.merged_result([result]) == result
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 random_allowed_result():
"""Return a RateLimitResult with is_allowed=True, otherwise random."""
return RateLimitResult(
is_allowed=True,
total_limit=randint(1, 100),
remaining_limit=randint(1, 100),
time_until_max=timedelta(randint(1, 100)),
)
def check_rate_limit(request: Request, action_name: str) -> RateLimitResult:
"""Check the rate limit for a particular action on a request."""
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_merged_all_allowed():
"""Ensure a merged result from all allowed results is also allowed."""
def random_allowed_result():
"""Return a RateLimitResult with is_allowed=True, otherwise random."""
return RateLimitResult(
is_allowed=True,
total_limit=randint(1, 100),
remaining_limit=randint(1, 100),
time_until_max=timedelta(randint(1, 100)),
)
# try merging a few different sets of different sizes
for num_results in range(2, 6):
results = [random_allowed_result() for _ in range(num_results)]
assert RateLimitResult.merged_result(results).is_allowed
def test_merged_results():
"""Ensure merging RateLimitResults gives the expected result."""
results = [
RateLimitResult(
is_allowed=True,
total_limit=20,
remaining_limit=15,
time_until_max=timedelta(seconds=90),
),
RateLimitResult(
is_allowed=False,
total_limit=10,
remaining_limit=0,
time_until_max=timedelta(seconds=30),
time_until_retry=timedelta(seconds=10),
),
RateLimitResult(
is_allowed=True,
total_limit=30,
remaining_limit=20,
time_until_max=timedelta(seconds=60),
),
]
expected_merged_result = RateLimitResult(
is_allowed=False,
total_limit=10,
remaining_limit=0,
time_until_max=timedelta(seconds=90),
time_until_retry=timedelta(seconds=10),
)
# try merging all permutations to ensure ordering isn't a factor
for permutation in permutations(results):
merged_result = RateLimitResult.merged_result(permutation)
assert merged_result == expected_merged_result