Skip to content

Commit 2620f75

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

3 files changed

Lines changed: 134 additions & 27 deletions

File tree

‎vulnerabilities/tests/test_api.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,9 +183,7 @@ def setUp(self):
183183
def test_package_serializer(self):
184184
pk = Package.objects.filter(name="mimetex").with_is_vulnerable()
185185
mock_request = RequestFactory().get("/api")
186-
response = PackageSerializer(
187-
pk, many=True, context={"request": mock_request}, HTTP_USER_AGENT="VCIO_API_AGENT"
188-
).data
186+
response = PackageSerializer(pk, many=True, context={"request": mock_request}).data
189187
self.assertEqual(1, len(response))
190188

191189
first_result = response[0]

‎vulnerabilities/tests/test_api_v3.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,20 @@ def setUp(self):
5656

5757
self.client = APIClient(enforce_csrf_checks=True)
5858

59+
self.allow_request_patcher = patch(
60+
"vulnerabilities.throttling.PermissionBasedUserRateThrottle.allow_request",
61+
return_value=True,
62+
)
63+
self.allow_request_patcher.start()
64+
self.addCleanup(self.allow_request_patcher.stop)
65+
66+
self.anon_patcher = patch(
67+
"rest_framework.throttling.AnonRateThrottle.allow_request",
68+
return_value=True,
69+
)
70+
self.anon_patcher.start()
71+
self.addCleanup(self.anon_patcher.stop)
72+
5973
def test_packages_post_without_details(self):
6074
url = reverse("package-v3-list")
6175

@@ -199,6 +213,19 @@ def setUp(self):
199213
advisory_obj.save()
200214

201215
self.client = APIClient(enforce_csrf_checks=True)
216+
self.allow_request_patcher = patch(
217+
"vulnerabilities.throttling.PermissionBasedUserRateThrottle.allow_request",
218+
return_value=True,
219+
)
220+
self.allow_request_patcher.start()
221+
self.addCleanup(self.allow_request_patcher.stop)
222+
223+
self.anon_patcher = patch(
224+
"rest_framework.throttling.AnonRateThrottle.allow_request",
225+
return_value=True,
226+
)
227+
self.anon_patcher.start()
228+
self.addCleanup(self.anon_patcher.stop)
202229

203230
def test_advisories_post(self):
204231
url = reverse("advisory-v3-list")
@@ -247,6 +274,20 @@ def setUp(self):
247274

248275
self.client = APIClient(enforce_csrf_checks=True)
249276

277+
self.allow_request_patcher = patch(
278+
"vulnerabilities.throttling.PermissionBasedUserRateThrottle.allow_request",
279+
return_value=True,
280+
)
281+
self.allow_request_patcher.start()
282+
self.addCleanup(self.allow_request_patcher.stop)
283+
284+
self.anon_patcher = patch(
285+
"rest_framework.throttling.AnonRateThrottle.allow_request",
286+
return_value=True,
287+
)
288+
self.anon_patcher.start()
289+
self.addCleanup(self.anon_patcher.stop)
290+
250291
def test_get_all_vulnerable_purls(self):
251292
url = reverse("package-v3-list")
252293

@@ -293,6 +334,20 @@ def setUp(self):
293334
url="https://github.com/aboutcode-org/sample",
294335
)
295336

337+
self.allow_request_patcher = patch(
338+
"vulnerabilities.throttling.PermissionBasedUserRateThrottle.allow_request",
339+
return_value=True,
340+
)
341+
self.allow_request_patcher.start()
342+
self.addCleanup(self.allow_request_patcher.stop)
343+
344+
self.anon_patcher = patch(
345+
"rest_framework.throttling.AnonRateThrottle.allow_request",
346+
return_value=True,
347+
)
348+
self.anon_patcher.start()
349+
self.addCleanup(self.anon_patcher.stop)
350+
296351
self.advisory = insert_advisory_v2(self.advisory, "importer_1", print, 100)
297352
self.advisory.is_latest = True
298353
self.advisory._all_impacts_unfurled_at = timezone.now()
@@ -356,6 +411,20 @@ def setUp(self):
356411
defaults={"name": "sample", "type": "pypi", "version": "1.0.0"},
357412
)
358413

414+
self.allow_request_patcher = patch(
415+
"vulnerabilities.throttling.PermissionBasedUserRateThrottle.allow_request",
416+
return_value=True,
417+
)
418+
self.allow_request_patcher.start()
419+
self.addCleanup(self.allow_request_patcher.stop)
420+
421+
self.anon_patcher = patch(
422+
"rest_framework.throttling.AnonRateThrottle.allow_request",
423+
return_value=True,
424+
)
425+
self.anon_patcher.start()
426+
self.addCleanup(self.anon_patcher.stop)
427+
359428
self.advisory_data = AdvisoryDataV2(
360429
advisory_id="AVID-123",
361430
aliases=[],
@@ -490,6 +559,21 @@ def test_advisory_set_member_patches_aggregation(self):
490559

491560

492561
class TestPackageTypesView(APITestCase):
562+
def setUp(self):
563+
self.allow_request_patcher = patch(
564+
"vulnerabilities.throttling.PermissionBasedUserRateThrottle.allow_request",
565+
return_value=True,
566+
)
567+
self.allow_request_patcher.start()
568+
self.addCleanup(self.allow_request_patcher.stop)
569+
570+
self.anon_patcher = patch(
571+
"rest_framework.throttling.AnonRateThrottle.allow_request",
572+
return_value=True,
573+
)
574+
self.anon_patcher.start()
575+
self.addCleanup(self.anon_patcher.stop)
576+
493577
def test_returns_distinct_types_and_caches_response(self):
494578
cache.delete("package_types")
495579

‎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)