unittests: import query/task count capture

docker/entrypoint-unit-tests-devDocker.sh

@@ -73,15 +73,15 @@ echo "Unit Tests"
 echo "------------------------------------------------------------"
 
 # Removing parallel and shuffle for now to maintain stability
-python3 manage.py test unittests -v 3 --keepdb --no-input --exclude-tag="non-parallel" || {
-    exit 1;
-}
-python3 manage.py test unittests -v 3 --keepdb --no-input --tag="non-parallel" || {
-    exit 1;
-}
+# python3 manage.py test unittests -v 3 --keepdb --no-input --exclude-tag="non-parallel" || {
+#     exit 1;
+# }
+# python3 manage.py test unittests -v 3 --keepdb --no-input --tag="non-parallel" || {
+#     exit 1;
+# }
 
 # you can select a single file to "test" unit tests
-# python3 manage.py test unittests.tools.test_npm_audit_scan_parser.TestNpmAuditParser --keepdb -v 3
+python3 manage.py test unittests.test_importers_performance.TestDojoImporterPerformance --keepdb -v 3 &> /app/dev2.log
 
 # or even a single method
 # python3 manage.py test unittests.tools.test_npm_audit_scan_parser.TestNpmAuditParser.test_npm_audit_parser_many_vuln_npm7 --keepdb -v 3

dojo/decorators.py

@@ -1,4 +1,5 @@
 import logging
+import threading
 from functools import wraps
 
 from django.conf import settings
@@ -12,6 +13,47 @@
 logger = logging.getLogger(__name__)
 
 
+class ThreadLocalTaskCounter:
+    def __init__(self):
+        self._thread_local = threading.local()
+
+    def _get_task_list(self):
+        if not hasattr(self._thread_local, "tasks"):
+            self._thread_local.tasks = []
+        return self._thread_local.tasks
+
+    def _get_recording(self):
+        return getattr(self._thread_local, "recording", False)
+
+    def start(self):
+        self._thread_local.recording = True
+        self._get_task_list().clear()
+
+    def stop(self):
+        self._thread_local.recording = False
+
+    def incr(self, task_name, model_id=None, args=None, kwargs=None):
+        if not self._get_recording():
+            return
+        tasks = self._get_task_list()
+        tasks.append({
+            "task": task_name,
+            "id": model_id,
+            "args": args if args is not None else [],
+            "kwargs": kwargs if kwargs is not None else {},
+        })
+
+    def get(self):
+        return len(self._get_task_list())
+
+    def get_tasks(self):
+        return list(self._get_task_list())
+
+
+# Create a shared instance
+dojo_async_task_counter = ThreadLocalTaskCounter()
+
+
 def we_want_async(*args, func=None, **kwargs):
     from dojo.models import Dojo_User
     from dojo.utils import get_current_user
@@ -40,6 +82,13 @@ def __wrapper__(*args, **kwargs):
         from dojo.utils import get_current_user
         user = get_current_user()
         kwargs["async_user"] = user
+
+        dojo_async_task_counter.incr(
+            func.__name__,
+            args=args,
+            kwargs=kwargs,
+        )
+
         countdown = kwargs.pop("countdown", 0)
         if we_want_async(*args, func=func, **kwargs):
             return func.apply_async(args=args, kwargs=kwargs, countdown=countdown)

unittests/dojo_test_case.py

@@ -97,14 +97,20 @@ def system_settings(
         enable_jira_web_hook=False,
         disable_jira_webhook_secret=False,
         jira_webhook_secret=None,
+        enable_github=False,
         enable_product_tag_inehritance=False,
+        enable_product_grade=True,
+        enable_webhooks_notifications=True,
     ):
         ss = System_Settings.objects.get()
         ss.enable_jira = enable_jira
         ss.enable_jira_web_hook = enable_jira_web_hook
         ss.disable_jira_webhook_secret = disable_jira_webhook_secret
         ss.jira_webhook_secret = jira_webhook_secret
+        ss.enable_github = enable_github
         ss.enable_product_tag_inheritance = enable_product_tag_inehritance
+        ss.enable_product_grade = enable_product_grade
+        ss.enable_webhooks_notifications = enable_webhooks_notifications
         ss.save()
 
     def create_product_type(self, name, *args, description="dummy description", **kwargs):