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