Skip to content

Commit ebfcaa7

Browse files
add foreground testcase
1 parent 9d6f4e5 commit ebfcaa7

1 file changed

Lines changed: 154 additions & 0 deletions

File tree

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
import logging
2+
from contextlib import contextmanager
3+
from unittest.mock import patch
4+
5+
from crum import impersonate
6+
from django.utils import timezone
7+
8+
from dojo.decorators import dojo_async_task_counter
9+
from dojo.importers.default_importer import DefaultImporter
10+
from dojo.importers.default_reimporter import DefaultReImporter
11+
from dojo.models import Development_Environment, Dojo_User, Engagement, Product, Product_Type, User
12+
13+
from .dojo_test_case import DojoTestCase, get_unit_tests_scans_path
14+
15+
logger = logging.getLogger(__name__)
16+
17+
STACK_HAWK_FILENAME = get_unit_tests_scans_path("stackhawk") / "stackhawk_many_vul_without_duplicated_findings.json"
18+
STACK_HAWK_SUBSET_FILENAME = get_unit_tests_scans_path("stackhawk") / "stackhawk_many_vul_without_duplicated_findings_subset.json"
19+
STACK_HAWK_SCAN_TYPE = "StackHawk HawkScan"
20+
21+
22+
class TestDojoImporterPerformance(DojoTestCase):
23+
fixtures = ["dojo_testdata.json"]
24+
25+
def setUp(self):
26+
super().setUp()
27+
user = User.objects.get(username="admin")
28+
user.usercontactinfo.block_execution = True
29+
user.save()
30+
31+
@contextmanager
32+
def assertNumAsyncTask(self, num):
33+
dojo_async_task_counter.start()
34+
try:
35+
yield
36+
finally:
37+
dojo_async_task_counter.stop()
38+
actual = dojo_async_task_counter.get()
39+
if actual != num:
40+
tasks = dojo_async_task_counter.get_tasks()
41+
tasks_str = "\n".join(str(task) for task in tasks)
42+
msg = (
43+
f"Expected {num} celery tasks, but {actual} were created.\n"
44+
f"Tasks created:\n{tasks_str}"
45+
)
46+
raise self.failureException(msg)
47+
48+
def import_reimport_performance(self, expected_num_queries1, expected_num_async_tasks1, expected_num_queries2, expected_num_async_tasks2, expected_num_queries3, expected_num_async_tasks3):
49+
"""Despite all efforts, these imports here run in async mode, so celery tasks are executed in the background"""
50+
product_type, _created = Product_Type.objects.get_or_create(name="test")
51+
product, _created = Product.objects.get_or_create(
52+
name="TestDojoDefaultImporter",
53+
prod_type=product_type,
54+
)
55+
engagement, _created = Engagement.objects.get_or_create(
56+
name="Test Create Engagement",
57+
product=product,
58+
target_start=timezone.now(),
59+
target_end=timezone.now(),
60+
)
61+
lead, _ = User.objects.get_or_create(username="admin")
62+
environment, _ = Development_Environment.objects.get_or_create(name="Development")
63+
64+
# first import the subset which missed one finding and a couple of endpoints on some of the findings
65+
with (
66+
self.subTest("import1"), impersonate(Dojo_User.objects.get(username="admin")),
67+
self.assertNumQueries(expected_num_queries1),
68+
self.assertNumAsyncTask(expected_num_async_tasks1),
69+
STACK_HAWK_SUBSET_FILENAME.open(encoding="utf-8") as scan,
70+
):
71+
import_options = {
72+
"user": lead,
73+
"lead": lead,
74+
"scan_date": None,
75+
"environment": environment,
76+
"minimum_severity": "Info",
77+
"active": True,
78+
"verified": True,
79+
"sync": True,
80+
"scan_type": STACK_HAWK_SCAN_TYPE,
81+
"engagement": engagement,
82+
}
83+
importer = DefaultImporter(**import_options)
84+
test, _, _len_new_findings, _len_closed_findings, _, _, _ = importer.process_scan(scan)
85+
86+
# use reimport with the full report so it add a finding and some endpoints
87+
with (
88+
self.subTest("reimport1"), impersonate(Dojo_User.objects.get(username="admin")),
89+
self.assertNumQueries(expected_num_queries2),
90+
self.assertNumAsyncTask(expected_num_async_tasks2),
91+
STACK_HAWK_FILENAME.open(encoding="utf-8") as scan,
92+
):
93+
reimport_options = {
94+
"test": test,
95+
"user": lead,
96+
"lead": lead,
97+
"scan_date": None,
98+
"minimum_severity": "Info",
99+
"active": True,
100+
"verified": True,
101+
"sync": True,
102+
"scan_type": STACK_HAWK_SCAN_TYPE,
103+
}
104+
reimporter = DefaultReImporter(**reimport_options)
105+
test, _, _len_new_findings, _len_closed_findings, _, _, _ = reimporter.process_scan(scan)
106+
107+
# use reimport with the subset again to close a finding and mitigate some endpoints
108+
with (
109+
self.subTest("reimport2"), impersonate(Dojo_User.objects.get(username="admin")),
110+
self.assertNumQueries(expected_num_queries3),
111+
self.assertNumAsyncTask(expected_num_async_tasks3),
112+
STACK_HAWK_SUBSET_FILENAME.open(encoding="utf-8") as scan,
113+
):
114+
reimport_options = {
115+
"test": test,
116+
"user": lead,
117+
"lead": lead,
118+
"scan_date": None,
119+
"minimum_severity": "Info",
120+
"active": True,
121+
"verified": True,
122+
"sync": True,
123+
"scan_type": STACK_HAWK_SCAN_TYPE,
124+
}
125+
reimporter = DefaultReImporter(**reimport_options)
126+
test, _, _len_new_findings, _len_closed_findings, _, _, _ = reimporter.process_scan(scan)
127+
128+
def test_import_reimport_reimport_performance(self):
129+
self.import_reimport_performance(
130+
expected_num_queries1=617,
131+
expected_num_async_tasks1=18,
132+
expected_num_queries2=496,
133+
expected_num_async_tasks2=25,
134+
expected_num_queries3=348,
135+
expected_num_async_tasks3=21,
136+
)
137+
138+
@patch("dojo.decorators.we_want_async", return_value=False)
139+
def test_import_reimport_reimport_performance_no_async(self, mock):
140+
"""
141+
This test checks the performance of the importers when they are run in sync mode.
142+
The reason for this is that we also want to be aware of when a PR affects the number of queries
143+
or async tasks created by a background task.
144+
The impersonate context manager above does not work as expected for disabling async,
145+
so we patch the we_want_async decorator to always return False.
146+
"""
147+
self.import_reimport_performance(
148+
expected_num_queries1=708,
149+
expected_num_async_tasks1=29,
150+
expected_num_queries2=566,
151+
expected_num_async_tasks2=32,
152+
expected_num_queries3=400,
153+
expected_num_async_tasks3=26,
154+
)

0 commit comments

Comments
 (0)