test_ratelimiting.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. from synapse.api.ratelimiting import Ratelimiter
  2. from tests import unittest
  3. class TestRatelimiter(unittest.TestCase):
  4. def test_allowed(self):
  5. limiter = Ratelimiter()
  6. allowed, time_allowed = limiter.send_message(
  7. user_id="test_id", time_now_s=0, msg_rate_hz=0.1, burst_count=1,
  8. )
  9. self.assertTrue(allowed)
  10. self.assertEquals(10., time_allowed)
  11. allowed, time_allowed = limiter.send_message(
  12. user_id="test_id", time_now_s=5, msg_rate_hz=0.1, burst_count=1,
  13. )
  14. self.assertFalse(allowed)
  15. self.assertEquals(10., time_allowed)
  16. allowed, time_allowed = limiter.send_message(
  17. user_id="test_id", time_now_s=10, msg_rate_hz=0.1, burst_count=1
  18. )
  19. self.assertTrue(allowed)
  20. self.assertEquals(20., time_allowed)
  21. def test_pruning(self):
  22. limiter = Ratelimiter()
  23. allowed, time_allowed = limiter.send_message(
  24. user_id="test_id_1", time_now_s=0, msg_rate_hz=0.1, burst_count=1,
  25. )
  26. self.assertIn("test_id_1", limiter.message_counts)
  27. allowed, time_allowed = limiter.send_message(
  28. user_id="test_id_2", time_now_s=10, msg_rate_hz=0.1, burst_count=1
  29. )
  30. self.assertNotIn("test_id_1", limiter.message_counts)