Skip to content

Commit cc9439e

Browse files
TG1999TG1999
committed
Fix tests
Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>
1 parent b0a3070 commit cc9439e

2 files changed

Lines changed: 54 additions & 25 deletions

File tree

‎vulnerabilities/tests/test_api_v2.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,11 @@ def test_list_packages(self):
245245
"""
246246
url = reverse("package-v2-list")
247247
with self.assertNumQueries(31):
248-
response = self.client.get(url, format="json", HTTP_USER_AGENT="VCIO_API_AGENT",)
248+
response = self.client.get(
249+
url,
250+
format="json",
251+
HTTP_USER_AGENT="VCIO_API_AGENT",
252+
)
249253
self.assertEqual(response.status_code, status.HTTP_200_OK)
250254
self.assertIn("results", response.data)
251255
self.assertIn("packages", response.data["results"])

‎vulnerabilities/tests/test_throttling.py‎

Lines changed: 49 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,10 @@
2222

2323
def simulate_throttle_usage(url, client, mock_use_count):
2424
throttle = PermissionBasedUserRateThrottle()
25-
request = client.get(url).wsgi_request
25+
26+
request = client.get(url, HTTP_USER_AGENT="VCIO_API_AGENT").wsgi_request
2627

2728
if cache_key := throttle.get_cache_key(request, view=None):
28-
print(cache_key)
2929
now = throttle.timer()
3030
cache.set(cache_key, [now] * mock_use_count)
3131

@@ -95,56 +95,72 @@ def test_user_with_low_perm_throttling(self):
9595
simulate_throttle_usage(
9696
url="/api/packages",
9797
client=self.th_low_user_csrf_client,
98-
mock_use_count=10799,
98+
mock_use_count=19,
9999
)
100100

101-
response = self.th_low_user_csrf_client.get("/api/packages")
101+
response = self.th_low_user_csrf_client.get(
102+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
103+
)
102104
self.assertEqual(response.status_code, status.HTTP_200_OK)
103105

104-
# exhausted 10800/hr allowed requests.
105-
response = self.th_low_user_csrf_client.get("/api/packages")
106+
# exhausted 20/minute allowed requests.
107+
response = self.th_low_user_csrf_client.get(
108+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
109+
)
106110
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
107111

108112
def test_basic_user_throttling(self):
109113
simulate_throttle_usage(
110114
url="/api/packages",
111115
client=self.basic_user_csrf_client,
112-
mock_use_count=14399,
116+
mock_use_count=29,
113117
)
114118

115-
response = self.basic_user_csrf_client.get("/api/packages")
119+
response = self.basic_user_csrf_client.get(
120+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
121+
)
116122
self.assertEqual(response.status_code, status.HTTP_200_OK)
117123

118-
# exhausted 14400/hr allowed requests.
119-
response = self.basic_user_csrf_client.get("/api/packages")
124+
# exhausted 30/minute allowed requests.
125+
response = self.basic_user_csrf_client.get(
126+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
127+
)
120128
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
121129

122130
def test_user_with_medium_perm_throttling(self):
123131
simulate_throttle_usage(
124132
url="/api/packages",
125133
client=self.th_medium_user_csrf_client,
126-
mock_use_count=14399,
134+
mock_use_count=29,
127135
)
128136

129-
response = self.th_medium_user_csrf_client.get("/api/packages")
137+
response = self.th_medium_user_csrf_client.get(
138+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
139+
)
130140
self.assertEqual(response.status_code, status.HTTP_200_OK)
131141

132-
# exhausted 14400/hr allowed requests for user with 14400 perm.
133-
response = self.th_medium_user_csrf_client.get("/api/packages")
142+
# exhausted 30/minute allowed requests for user with 14400 perm.
143+
response = self.th_medium_user_csrf_client.get(
144+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
145+
)
134146
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
135147

136148
def test_user_with_high_perm_throttling(self):
137149
simulate_throttle_usage(
138150
url="/api/packages",
139151
client=self.th_high_user_csrf_client,
140-
mock_use_count=17999,
152+
mock_use_count=0,
141153
)
142154

143-
response = self.th_high_user_csrf_client.get("/api/packages")
155+
response = self.th_high_user_csrf_client.get(
156+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
157+
)
144158
self.assertEqual(response.status_code, status.HTTP_200_OK)
145159

146160
# exhausted 18000/hr allowed requests for user with 18000 perm.
147-
response = self.th_high_user_csrf_client.get("/api/packages")
161+
response = self.th_high_user_csrf_client.get(
162+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
163+
)
148164
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
149165

150166
def test_user_with_unrestricted_perm_throttling(self):
@@ -155,7 +171,9 @@ def test_user_with_unrestricted_perm_throttling(self):
155171
)
156172

157173
# no throttling for user with unrestricted perm.
158-
response = self.th_unrestricted_user_csrf_client.get("/api/packages")
174+
response = self.th_unrestricted_user_csrf_client.get(
175+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
176+
)
159177
self.assertEqual(response.status_code, status.HTTP_200_OK)
160178

161179
def test_user_in_group_with_unrestricted_perm_throttling(self):
@@ -166,28 +184,32 @@ def test_user_in_group_with_unrestricted_perm_throttling(self):
166184
)
167185

168186
# no throttling for user in group with unrestricted perm.
169-
response = self.th_group_user_csrf_client.get("/api/packages")
187+
response = self.th_group_user_csrf_client.get(
188+
"/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT"
189+
)
170190
self.assertEqual(response.status_code, status.HTTP_200_OK)
171191

172192
def test_anon_throttling(self):
173193
simulate_throttle_usage(
174194
url="/api/packages",
175195
client=self.csrf_client_anon,
176-
mock_use_count=3599,
196+
mock_use_count=9,
177197
)
178198

179-
response = self.csrf_client_anon.get("/api/packages")
199+
response = self.csrf_client_anon.get("/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT")
180200
self.assertEqual(response.status_code, status.HTTP_200_OK)
181201

182202
# exhausted 3600/hr allowed requests for anon.
183-
response = self.csrf_client_anon.get("/api/packages")
203+
response = self.csrf_client_anon.get("/api/packages", HTTP_USER_AGENT="VCIO_API_AGENT")
184204
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
185205
self.assertEqual(
186206
response.data.get("message"),
187207
"Your request has been throttled. Please contact support@nexb.com",
188208
)
189209

190-
response = self.csrf_client_anon.get("/api/vulnerabilities")
210+
response = self.csrf_client_anon.get(
211+
"/api/vulnerabilities", HTTP_USER_AGENT="VCIO_API_AGENT"
212+
)
191213
# 429 - too many requests for anon user
192214
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)
193215
self.assertEqual(
@@ -198,7 +220,10 @@ def test_anon_throttling(self):
198220
data = json.dumps({"purls": ["pkg:foo/bar"]})
199221

200222
response = self.csrf_client_anon.post(
201-
"/api/packages/bulk_search", data=data, content_type="application/json"
223+
"/api/packages/bulk_search",
224+
data=data,
225+
content_type="application/json",
226+
HTTP_USER_AGENT="VCIO_API_AGENT",
202227
)
203228
# 429 - too many requests for anon user
204229
self.assertEqual(response.status_code, status.HTTP_429_TOO_MANY_REQUESTS)

0 commit comments

Comments
 (0)