Skip to content

Commit 71fd383

Browse files
shtlrssupakeen
authored andcommitted
add tests to cover the ratelimiting behavior
1 parent 5e0376f commit 71fd383

File tree

1 file changed

+90
-0
lines changed

1 file changed

+90
-0
lines changed
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import copy
2+
import time
3+
import unittest.mock
4+
5+
import tornado.testing
6+
import tornado.web
7+
8+
from pinnwand import app, configuration, defensive
9+
10+
11+
class RateLimitTestCase(tornado.testing.AsyncHTTPTestCase):
12+
13+
def get_app(self) -> tornado.web.Application:
14+
return app.make_application()
15+
16+
def test_ratelimit_verification_on_endpoints(self):
17+
with unittest.mock.patch("pinnwand.defensive.should_be_ratelimited") as patch:
18+
patch.return_value = False
19+
20+
self.fetch(
21+
"/",
22+
method="GET",
23+
)
24+
25+
patch.assert_called()
26+
patch.reset_mock()
27+
28+
def test_ratelimit_application_on_one_client(self):
29+
config = configuration.ConfigurationProvider.get_config()
30+
ratelimlit_copy = copy.deepcopy(config._ratelimit)
31+
ratelimlit_copy["read"]["capacity"] = 2
32+
ratelimlit_copy["read"]["consume"] = 2
33+
ratelimlit_copy["read"]["refill"] = 1
34+
35+
with unittest.mock.patch.dict("pinnwand.defensive.ConfigurationProvider._config._ratelimit", ratelimlit_copy):
36+
with unittest.mock.patch.dict("pinnwand.defensive.ratelimit_area", clear=True):
37+
response = self.fetch(
38+
"/",
39+
method="GET",
40+
)
41+
42+
assert response.code == 200
43+
44+
response = self.fetch(
45+
"/",
46+
method="GET",
47+
)
48+
49+
assert response.code == 429
50+
51+
def test_ratelimit_application_on_multiple_clients(self):
52+
config = configuration.ConfigurationProvider.get_config()
53+
ratelimlit_copy = copy.deepcopy(config._ratelimit)
54+
area = "read"
55+
ratelimlit_copy[area]["capacity"] = 10
56+
ratelimlit_copy[area]["consume"] = 7
57+
ratelimlit_copy[area]["refill"] = 1
58+
59+
ip1 = "192.168.15.32"
60+
ip2 = "10.45.134.23"
61+
62+
with unittest.mock.patch.dict("pinnwand.defensive.ConfigurationProvider._config._ratelimit", ratelimlit_copy):
63+
with unittest.mock.patch.dict("pinnwand.defensive.ratelimit_area", clear=True):
64+
assert defensive.should_be_ratelimited(ip1, area) is False
65+
assert defensive.should_be_ratelimited(ip1, area) is True
66+
assert defensive.should_be_ratelimited(ip2, area) is False
67+
assert defensive.should_be_ratelimited(ip2, area) is True
68+
assert defensive.should_be_ratelimited(ip2, area) is True
69+
time.sleep(10) # Give it enough time to replenish
70+
assert defensive.should_be_ratelimited(ip1, area) is False
71+
assert defensive.should_be_ratelimited(ip2, area) is False
72+
73+
def test_bucket_tokens_consumption(self):
74+
config = configuration.ConfigurationProvider.get_config()
75+
ratelimlit_copy = copy.deepcopy(config._ratelimit)
76+
area = "read"
77+
consumption = 7
78+
capacity = 10
79+
ratelimlit_copy[area]["capacity"] = capacity
80+
ratelimlit_copy[area]["consume"] = consumption
81+
ratelimlit_copy[area]["refill"] = 1
82+
83+
ip = "192.168.15.32"
84+
with unittest.mock.patch.dict("pinnwand.defensive.ConfigurationProvider._config._ratelimit", ratelimlit_copy):
85+
with unittest.mock.patch.dict("pinnwand.defensive.ratelimit_area", clear=True):
86+
defensive.should_be_ratelimited(ip, area)
87+
limiter = defensive.ratelimit_area[area]
88+
tokens_remaining = limiter._storage.get_token_count(ip.encode("utf-8"))
89+
assert tokens_remaining == capacity - consumption
90+

0 commit comments

Comments
 (0)