Add type hints to TestRatelimiter (#14885)

This commit is contained in:
Andrew Morgan 2023-01-21 15:59:15 +00:00 committed by GitHub
parent 0ec12a3753
commit 8d90e5f200
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 52 additions and 16 deletions

1
changelog.d/14885.misc Normal file
View File

@ -0,0 +1 @@
Add missing type hints.

View File

@ -33,7 +33,6 @@ exclude = (?x)
|synapse/storage/schema/ |synapse/storage/schema/
|tests/api/test_auth.py |tests/api/test_auth.py
|tests/api/test_ratelimiting.py
|tests/app/test_openid_listener.py |tests/app/test_openid_listener.py
|tests/appservice/test_scheduler.py |tests/appservice/test_scheduler.py
|tests/events/test_presence_router.py |tests/events/test_presence_router.py

View File

@ -8,7 +8,10 @@ from tests import unittest
class TestRatelimiter(unittest.HomeserverTestCase): class TestRatelimiter(unittest.HomeserverTestCase):
def test_allowed_via_can_do_action(self): def test_allowed_via_can_do_action(self):
limiter = Ratelimiter( 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( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(None, key="test_id", _time_now_s=0) limiter.can_do_action(None, key="test_id", _time_now_s=0)
@ -30,7 +33,7 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_allowed_appservice_ratelimited_via_can_requester_do_action(self): def test_allowed_appservice_ratelimited_via_can_requester_do_action(self):
appservice = ApplicationService( appservice = ApplicationService(
None, token="fake_token",
id="foo", id="foo",
rate_limited=True, rate_limited=True,
sender="@as:example.com", sender="@as:example.com",
@ -38,7 +41,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
as_requester = create_requester("@user:example.com", app_service=appservice) as_requester = create_requester("@user:example.com", app_service=appservice)
limiter = Ratelimiter( 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( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(as_requester, _time_now_s=0) limiter.can_do_action(as_requester, _time_now_s=0)
@ -60,7 +66,7 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_allowed_appservice_via_can_requester_do_action(self): def test_allowed_appservice_via_can_requester_do_action(self):
appservice = ApplicationService( appservice = ApplicationService(
None, token="fake_token",
id="foo", id="foo",
rate_limited=False, rate_limited=False,
sender="@as:example.com", sender="@as:example.com",
@ -68,7 +74,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
as_requester = create_requester("@user:example.com", app_service=appservice) as_requester = create_requester("@user:example.com", app_service=appservice)
limiter = Ratelimiter( 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( allowed, time_allowed = self.get_success_or_raise(
limiter.can_do_action(as_requester, _time_now_s=0) limiter.can_do_action(as_requester, _time_now_s=0)
@ -90,7 +99,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_allowed_via_ratelimit(self): def test_allowed_via_ratelimit(self):
limiter = Ratelimiter( 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 # Shouldn't raise
@ -114,7 +126,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
""" """
# Create a Ratelimiter with a very low allowed rate_hz and burst_count # Create a Ratelimiter with a very low allowed rate_hz and burst_count
limiter = Ratelimiter( 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 # First attempt should be allowed
@ -160,7 +175,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
""" """
# Create a Ratelimiter with a very low allowed rate_hz and burst_count # Create a Ratelimiter with a very low allowed rate_hz and burst_count
limiter = Ratelimiter( 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 # First attempt should be allowed
@ -188,7 +206,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_pruning(self): def test_pruning(self):
limiter = Ratelimiter( 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( self.get_success_or_raise(
limiter.can_do_action(None, key="test_id_1", _time_now_s=0) limiter.can_do_action(None, key="test_id_1", _time_now_s=0)
@ -223,7 +244,7 @@ class TestRatelimiter(unittest.HomeserverTestCase):
) )
) )
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 # Shouldn't raise
for _ in range(20): for _ in range(20):
@ -231,7 +252,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_multiple_actions(self): def test_multiple_actions(self):
limiter = Ratelimiter( 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. # Test that 4 actions aren't allowed with a maximum burst of 3.
allowed, time_allowed = self.get_success_or_raise( allowed, time_allowed = self.get_success_or_raise(
@ -295,7 +319,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
extra tokens by timing requests. extra tokens by timing requests.
""" """
limiter = Ratelimiter( 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: def consume_at(time: float) -> bool:
@ -317,7 +344,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_record_action_which_doesnt_fill_bucket(self) -> None: def test_record_action_which_doesnt_fill_bucket(self) -> None:
limiter = Ratelimiter( 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. # Observe two actions, leaving room in the bucket for one more.
@ -337,7 +367,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_record_action_which_fills_bucket(self) -> None: def test_record_action_which_fills_bucket(self) -> None:
limiter = Ratelimiter( 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. # Observe three actions, filling up the bucket.
@ -363,7 +396,10 @@ class TestRatelimiter(unittest.HomeserverTestCase):
def test_record_action_which_overfills_bucket(self) -> None: def test_record_action_which_overfills_bucket(self) -> None:
limiter = Ratelimiter( 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. # Observe four actions, exceeding the bucket.