Skip to content

Commit

Permalink
Add type hints to TestRatelimiter (matrix-org#14885)
Browse files Browse the repository at this point in the history
  • Loading branch information
anoadragon453 authored Jan 21, 2023
1 parent 0ec12a3 commit 8d90e5f
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 16 deletions.
1 change: 1 addition & 0 deletions changelog.d/14885.misc
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add missing type hints.
1 change: 0 additions & 1 deletion mypy.ini
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ exclude = (?x)
|synapse/storage/schema/

|tests/api/test_auth.py
|tests/api/test_ratelimiting.py
|tests/app/test_openid_listener.py
|tests/appservice/test_scheduler.py
|tests/events/test_presence_router.py
Expand Down
66 changes: 51 additions & 15 deletions tests/api/test_ratelimiting.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,10 @@
class TestRatelimiter(unittest.HomeserverTestCase):
def test_allowed_via_can_do_action(self):
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)
allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(None, key="test_id", _time_now_s=0)
Expand All @@ -30,15 +33,18 @@ def test_allowed_via_can_do_action(self):

def test_allowed_appservice_ratelimited_via_can_requester_do_action(self):
appservice = ApplicationService(
None,
token="fake_token",
id="foo",
rate_limited=True,
sender="@as:example.com",
)
as_requester = create_requester("@user:example.com", app_service=appservice)

limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)
allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(as_requester, _time_now_s=0)
Expand All @@ -60,15 +66,18 @@ def test_allowed_appservice_ratelimited_via_can_requester_do_action(self):

def test_allowed_appservice_via_can_requester_do_action(self):
appservice = ApplicationService(
None,
token="fake_token",
id="foo",
rate_limited=False,
sender="@as:example.com",
)
as_requester = create_requester("@user:example.com", app_service=appservice)

limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)
allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(as_requester, _time_now_s=0)
Expand All @@ -90,7 +99,10 @@ def test_allowed_appservice_via_can_requester_do_action(self):

def test_allowed_via_ratelimit(self):
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)

# Shouldn't raise
Expand All @@ -114,7 +126,10 @@ def test_allowed_via_can_do_action_and_overriding_parameters(self):
"""
# Create a Ratelimiter with a very low allowed rate_hz and burst_count
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)

# First attempt should be allowed
Expand Down Expand Up @@ -160,7 +175,10 @@ def test_allowed_via_ratelimit_and_overriding_parameters(self):
"""
# Create a Ratelimiter with a very low allowed rate_hz and burst_count
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)

# First attempt should be allowed
Expand Down Expand Up @@ -188,7 +206,10 @@ def test_allowed_via_ratelimit_and_overriding_parameters(self):

def test_pruning(self):
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=1,
)
self.get_success_or_raise(
limiter.can_do_action(None, key="test_id_1", _time_now_s=0)
Expand Down Expand Up @@ -223,15 +244,18 @@ def test_db_user_override(self):
)
)

limiter = Ratelimiter(store=store, clock=None, rate_hz=0.1, burst_count=1)
limiter = Ratelimiter(store=store, clock=self.clock, rate_hz=0.1, burst_count=1)

# Shouldn't raise
for _ in range(20):
self.get_success_or_raise(limiter.ratelimit(requester, _time_now_s=0))

def test_multiple_actions(self):
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=3,
)
# Test that 4 actions aren't allowed with a maximum burst of 3.
allowed, time_allowed = self.get_success_or_raise(
Expand Down Expand Up @@ -295,7 +319,10 @@ def test_rate_limit_burst_only_given_once(self) -> None:
extra tokens by timing requests.
"""
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=3,
)

def consume_at(time: float) -> bool:
Expand All @@ -317,7 +344,10 @@ def consume_at(time: float) -> bool:

def test_record_action_which_doesnt_fill_bucket(self) -> None:
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=3,
)

# Observe two actions, leaving room in the bucket for one more.
Expand All @@ -337,7 +367,10 @@ def test_record_action_which_doesnt_fill_bucket(self) -> None:

def test_record_action_which_fills_bucket(self) -> None:
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=3,
)

# Observe three actions, filling up the bucket.
Expand All @@ -363,7 +396,10 @@ def test_record_action_which_fills_bucket(self) -> None:

def test_record_action_which_overfills_bucket(self) -> None:
limiter = Ratelimiter(
store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
store=self.hs.get_datastores().main,
clock=self.clock,
rate_hz=0.1,
burst_count=3,
)

# Observe four actions, exceeding the bucket.
Expand Down

0 comments on commit 8d90e5f

Please sign in to comment.