def test_override_rate_limit(self):
"""Override one or more of the default rate limits"""
class TestEndpoint(Endpoint):
rate_limits = {
"GET": {
RateLimitCategory.IP: RateLimit(100, 5)
},
"POST": {
RateLimitCategory.USER: RateLimit(20, 4)
},
}
assert get_rate_limit_value("GET", TestEndpoint,
"ip") == RateLimit(100, 5)
# get is not overriddent for user, hence we use the default
assert get_rate_limit_value(
"GET", TestEndpoint, "user") == get_default_rate_limits_for_group(
"default", category=RateLimitCategory.USER)
# get is not overriddent for IP, hence we use the default
assert get_rate_limit_value("POST", TestEndpoint,
"ip") == get_default_rate_limits_for_group(
"default",
category=RateLimitCategory.IP)
assert get_rate_limit_value("POST", TestEndpoint,
"user") == RateLimit(20, 4)
def test_default_rate_limit_values(self):
"""Ensure that the default rate limits are called for endpoints without overrides"""
class TestEndpoint(Endpoint):
pass
assert get_rate_limit_value("GET", TestEndpoint, "ip") == get_default_rate_limits_for_group(
"default", RateLimitCategory.IP
)
assert get_rate_limit_value(
"POST", TestEndpoint, "org"
) == get_default_rate_limits_for_group("default", RateLimitCategory.ORGANIZATION)
assert get_rate_limit_value(
"DELETE", TestEndpoint, "user"
) == get_default_rate_limits_for_group("default", RateLimitCategory.USER)
def test_override(self):
config = RateLimitConfig(
group="default",
limit_overrides={"GET": {
RateLimitCategory.IP: RateLimit(1, 1)
}})
assert config.get_rate_limit("GET",
RateLimitCategory.IP) == RateLimit(1, 1)
assert config.get_rate_limit(
"POST", RateLimitCategory.IP) == get_default_rate_limits_for_group(
"default", RateLimitCategory.IP)
assert config.get_rate_limit("GET", RateLimitCategory.ORGANIZATION
) == get_default_rate_limits_for_group(
"default",
RateLimitCategory.ORGANIZATION)
def test_invalid_config(self):
config = RateLimitConfig(
group="default", limit_overrides={"GET": {
"invalid": "invalid"
}}) # type: ignore
assert config.get_rate_limit(
"bloop", "badcategory") == get_default_rate_limits_for_group(
"default", RateLimitCategory.ORGANIZATION)
def test_override_rate_limit(self):
"""Override one or more of the default rate limits."""
class TestEndpoint(Endpoint):
rate_limits = {
"GET": {RateLimitCategory.IP: RateLimit(100, 5)},
"POST": {RateLimitCategory.USER: RateLimit(20, 4)},
}
assert get_rate_limit_value("GET", TestEndpoint, RateLimitCategory.IP) == RateLimit(100, 5)
assert get_rate_limit_value(
"GET", TestEndpoint, RateLimitCategory.USER
) == get_default_rate_limits_for_group("default", RateLimitCategory.USER)
assert get_rate_limit_value(
"POST", TestEndpoint, RateLimitCategory.IP
) == get_default_rate_limits_for_group("default", RateLimitCategory.IP)
assert get_rate_limit_value("POST", TestEndpoint, RateLimitCategory.USER) == RateLimit(
20, 4
)
def test_inherit(self):
class ParentEndpoint(Endpoint):
rate_limits = RateLimitConfig(
group="foo", limit_overrides={"GET": {RateLimitCategory.IP: RateLimit(100, 5)}}
)
class ChildEndpoint(ParentEndpoint):
rate_limits = RateLimitConfig(group="foo", limit_overrides={"GET": {}})
assert get_rate_limit_value(
"GET", ChildEndpoint, RateLimitCategory.IP
) == get_default_rate_limits_for_group("foo", RateLimitCategory.IP)