def test_apply_quota(
enabled: int,
referrer: str,
config_to_set: str,
expected_quota: Optional[ResourceQuota],
) -> None:
state.set_config(ENABLED_CONFIG, enabled)
state.set_config(config_to_set, 5)
query = Query(
QueryEntity(EntityKey.EVENTS, EntityColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=binary_condition(
ConditionFunctions.EQ,
Column("_snuba_project_id", None, "project_id"),
Literal(None, 1),
),
)
settings = HTTPQuerySettings()
settings.referrer = referrer
ResourceQuotaProcessor("project_id").process_query(query, settings)
assert settings.get_resource_quota() == expected_quota
def test_referrer_rate_limit_processor_no_config(
unprocessed: Expression, project_id: int
) -> None:
query = Query(
QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[SelectedExpression("column2", Column(None, None, "column2"))],
condition=unprocessed,
)
# don't configure it, rate limit shouldn't fire
referrer = "abusive_delivery"
settings = HTTPQuerySettings()
settings.referrer = referrer
num_before = len(settings.get_rate_limit_params())
ProjectReferrerRateLimiter("project_id").process_query(query, settings)
assert len(settings.get_rate_limit_params()) == num_before
def test_referrer_specified_project(unprocessed: Expression, project_id: int) -> None:
query = Query(
QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[SelectedExpression("column2", Column(None, None, "column2"))],
condition=unprocessed,
)
state.set_config("project_referrer_per_second_limit_abusive_delivery_1", 10)
state.set_config("project_referrer_concurrent_limit_abusive_delivery_1", 10)
referrer = "abusive_delivery"
settings = HTTPQuerySettings()
settings.referrer = referrer
num_before = len(settings.get_rate_limit_params())
ProjectReferrerRateLimiter("project_id").process_query(query, settings)
assert len(settings.get_rate_limit_params()) == num_before + 1
rate_limiter = settings.get_rate_limit_params()[-1]
assert rate_limiter.rate_limit_name == PROJECT_REFERRER_RATE_LIMIT_NAME
assert rate_limiter.bucket == f"{project_id}"
assert rate_limiter.per_second_limit == 10
assert rate_limiter.concurrent_limit == 10