Refine the way Notification are created and add tests #106

Signed-off-by: tdruez <[email protected]>

dejacode/settings.py

@@ -650,17 +650,14 @@ def get_fake_redis_connection(config, use_strict_redis):
     "request_comment.added": "workflow.RequestComment.created+",
     "user.added_or_updated": None,
     "user.locked_out": None,
-    "vulnerability.added": "vulnerabilities.Vulnerability.created+",
-    "vulnerability.packages_affected": None,
-    "vulnerability.products_affected": None,
+    # "vulnerability.added": "vulnerabilities.Vulnerability.created+",
+    # "vulnerability.packages_affected": None,
+    # "vulnerability.products_affected": None,
     "vulnerability.data_update": None,
 }
 # Provide context variables to the `Webhook` values such as `extra_headers`.
 HOOK_ENV = env.dict("HOOK_ENV", default={})
 
-# Internal notifications
-DJANGO_NOTIFICATIONS_CONFIG = {"USE_JSONFIELD": True}
-
 # Django-axes
 # Enable or disable Axes plugin functionality
 AXES_ENABLED = env.bool("AXES_ENABLED", default=False)

dje/models.py

@@ -1757,6 +1757,14 @@ def serialize_hook(self, hook):
             **self.serialize_user_data(),
         }
 
+    def internal_notify(self, verb, **data):
+        """Similar to ``notify.send`` without relying on the Signal system."""
+        return Notification.objects.create(
+            recipient=self,
+            verb=verb,
+            **data,
+        )
+
 
 @receiver(models.signals.post_save, sender=settings.AUTH_USER_MODEL)
 def create_auth_token(sender, instance=None, created=False, **kwargs):

dje/tests/test_notification.py

@@ -72,6 +72,13 @@ def test_get_data_update_recipients(self):
             email_list, get_user_model().objects.get_data_update_recipients(self.dataspace)
         )
 
+    def test_get_vulnerability_notifications_users(self):
+        qs = get_user_model().objects.get_vulnerability_notifications_users(self.dataspace)
+        self.assertEqual(0, qs.count())
+        self.user.vulnerability_impact_notification = True
+        self.user.save()
+        self.assertEqual(1, qs.count())
+
     def test_send_notification_email(self):
         request = self.factory.get("")
         request.user = self.user

vulnerabilities/fetch.py

@@ -8,18 +8,16 @@
 
 from timeit import default_timer as timer
 
-from django.contrib.auth import get_user_model
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.humanize.templatetags.humanize import intcomma
 from django.core.management.base import CommandError
 from django.urls import reverse
 from django.utils import timezone
 
-from notifications.signals import notify
-
 from component_catalog.models import PACKAGE_URL_FIELDS
 from component_catalog.models import Package
 from dejacode_toolkit.vulnerablecode import VulnerableCode
+from dje.models import DejacodeUser
 from dje.utils import chunked_queryset
 from dje.utils import humanize_time
 from notification.models import find_and_fire_hook
@@ -185,11 +183,11 @@ def notify_vulnerability_data_update(dataspace):
     )
 
     # 2. Internal notifications
-    users_to_notify = get_user_model().objects.get_vulnerability_notifications_users(dataspace)
-    notify.send(
-        sender=Vulnerability,
-        verb="New vulnerabilities detected",
-        recipient=users_to_notify,
-        description=f"{message}",
-        action_object_content_type=ContentType.objects.get_for_model(Vulnerability),
-    )
+    users_to_notify = DejacodeUser.objects.get_vulnerability_notifications_users(dataspace)
+    for user in users_to_notify:
+        user.internal_notify(
+            verb="New vulnerabilities detected",
+            description=f"{message}",
+            actor_content_type=ContentType.objects.get_for_model(Vulnerability),
+            action_object_content_type=ContentType.objects.get_for_model(Vulnerability),
+        )

vulnerabilities/tests/test_fetch.py

@@ -14,14 +14,18 @@
 
 from django.test import TestCase
 
+from notifications.models import Notification
+
 from component_catalog.models import Package
 from component_catalog.tests import make_package
 from dje.models import Dataspace
+from dje.tests import create_user
 from product_portfolio.tests import make_product
 from product_portfolio.tests import make_product_item_purpose
 from product_portfolio.tests import make_product_package
 from vulnerabilities.fetch import fetch_for_packages
 from vulnerabilities.fetch import fetch_from_vulnerablecode
+from vulnerabilities.fetch import notify_vulnerability_data_update
 
 
 class VulnerabilitiesFetchTestCase(TestCase):
@@ -96,3 +100,15 @@ def test_vulnerabilities_fetch_for_packages(self, mock_bulk_search_by_purl):
         pp1.refresh_from_db()
         self.assertEqual(Decimal("8.4"), package1.risk_score)
         self.assertEqual(Decimal("4.2"), pp1.weighted_risk_score)
+
+    @mock.patch("vulnerabilities.fetch.find_and_fire_hook")
+    def test_vulnerabilities_fetch_notify_vulnerability_data_update(self, mock_fire_hook):
+        notify_vulnerability_data_update(self.dataspace)
+        mock_fire_hook.assert_not_called()
+        self.assertEqual(0, Notification.objects.count())
+
+        make_package(self.dataspace, is_vulnerable=True)
+        create_user("test", self.dataspace, vulnerability_impact_notification=True)
+        notify_vulnerability_data_update(self.dataspace)
+        mock_fire_hook.assert_called_once()
+        self.assertEqual(1, Notification.objects.count())