def test_table_rate_limit(query: Query, limit_to_set: str,
params: RateLimitParameters) -> None:
set_config(limit_to_set, 50)
request_settings = HTTPRequestSettings(consistent=True)
TableRateLimit().process_query(query, request_settings)
rate_limiters = request_settings.get_rate_limit_params()
assert params in rate_limiters
def test_project_extension_query_adds_rate_limits() -> None:
extension = ProjectExtension(project_column="project_id")
raw_data = {"project": [1, 2]}
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({"conditions": []}, QueryEntity(EntityKey.EVENTS, ColumnSet([])))
request_settings = HTTPRequestSettings()
num_rate_limits_before_processing = len(request_settings.get_rate_limit_params())
extension.get_processor().process_query(query, valid_data, request_settings)
rate_limits = request_settings.get_rate_limit_params()
# make sure a rate limit was added by the processing
assert len(rate_limits) == num_rate_limits_before_processing + 1
most_recent_rate_limit = rate_limits[-1]
assert most_recent_rate_limit.bucket == "1"
assert most_recent_rate_limit.per_second_limit == 1000
assert most_recent_rate_limit.concurrent_limit == 1000
def test_project_rate_limit_processor(unprocessed: Expression,
project_id: int) -> None:
query = Query(
QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=unprocessed,
)
settings = HTTPRequestSettings()
num_before = len(settings.get_rate_limit_params())
ProjectRateLimiterProcessor("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_RATE_LIMIT_NAME
assert rate_limiter.bucket == str(project_id)
assert rate_limiter.per_second_limit == 1000
assert rate_limiter.concurrent_limit == 1000
def test_project_extension_project_rate_limits_are_overridden() -> None:
extension = ProjectExtension(project_column="project_id")
raw_data = {"project": [3, 4]}
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({"conditions": []}, QueryEntity(EntityKey.EVENTS, ColumnSet([])))
request_settings = HTTPRequestSettings()
state.set_config("project_per_second_limit_3", 5)
state.set_config("project_concurrent_limit_3", 10)
extension.get_processor().process_query(query, valid_data, request_settings)
rate_limits = request_settings.get_rate_limit_params()
most_recent_rate_limit = rate_limits[-1]
assert most_recent_rate_limit.bucket == "3"
assert most_recent_rate_limit.per_second_limit == 5
assert most_recent_rate_limit.concurrent_limit == 10
def test_project_extension_project_rate_limits_are_overridden():
extension = ProjectExtension(
processor=ProjectExtensionProcessor(project_column="project_id")
)
raw_data = {"project": [2, 3]}
valid_data = validate_jsonschema(raw_data, extension.get_schema())
query = Query({"conditions": []}, TableSource("my_table", ColumnSet([])),)
request_settings = HTTPRequestSettings()
state.set_config("project_per_second_limit_2", 5)
state.set_config("project_concurrent_limit_2", 10)
extension.get_processor().process_query(query, valid_data, request_settings)
rate_limits = request_settings.get_rate_limit_params()
most_recent_rate_limit = rate_limits[-1]
assert most_recent_rate_limit.bucket == "2"
assert most_recent_rate_limit.per_second_limit == 5
assert most_recent_rate_limit.concurrent_limit == 10
