Add proper method to update and set the risk_score

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

Author: tdruez

product_portfolio/models.py

@@ -756,7 +756,7 @@ def update_license_unknown(self):
             product_package.update_license_unknown()
 
     def annotate_weighted_risk_score(self):
-        """Annotate the Queeryset with the weighted_risk_score computed value."""
+        """Annotate the Queryset with the weighted_risk_score computed value."""
         purpose = ProductItemPurpose.objects.filter(productpackage=OuterRef("pk"))
         package = Package.objects.filter(productpackages=OuterRef("pk"))
 

vulnerabilities/models.py

@@ -428,6 +428,15 @@ class Meta:
     def is_vulnerable(self):
         return self.affected_by_vulnerabilities.exists()
 
+    def update_risk_score(self):
+        """Calculate and save the maximum risk score from all affected vulnerabilities."""
+        qs = self.affected_by_vulnerabilities.aggregate(models.Max("risk_score"))
+        max_score = qs["risk_score__max"]
+
+        self.risk_score = max_score
+        self.save(update_fields=["risk_score"])
+        return self.risk_score
+
     def get_entry_for_package(self, vulnerablecode):
         if not self.package_url:
             return
@@ -495,8 +504,7 @@ def create_vulnerabilities(self, vulnerabilities_data):
         through_defaults = {"dataspace_id": self.dataspace_id}
         self.affected_by_vulnerabilities.add(*vulnerabilities, through_defaults=through_defaults)
 
-        # TODO: Looks like a bug....
-        self.update(risk_score=vulnerability_data["risk_score"])
+        self.update_risk_score()
         if isinstance(self, Package):
             self.productpackages.update_weighted_risk_score()
 

vulnerabilities/tests/test_models.py

@@ -82,19 +82,50 @@ def test_vulnerability_mixin_create_vulnerabilities(self):
         response_file = self.data / "vulnerabilities" / "idna_3.6_response.json"
         response_json = json.loads(response_file.read_text())
         vulnerabilities_data = response_json["results"][0]["affected_by_vulnerabilities"]
+        vulnerabilities_data.append({"vulnerability_id": "VCID-0002", "risk_score": 5.0})
 
         package1 = make_package(self.dataspace, package_url="pkg:pypi/[email protected]")
         product1 = make_product(self.dataspace, inventory=[package1])
         package1.create_vulnerabilities(vulnerabilities_data)
 
-        self.assertEqual(1, Vulnerability.objects.scope(self.dataspace).count())
-        self.assertEqual(1, package1.affected_by_vulnerabilities.count())
-        vulnerability = package1.affected_by_vulnerabilities.get()
-        self.assertEqual("VCID-j3au-usaz-aaag", vulnerability.vulnerability_id)
-
-        self.assertEqual(8.4, package1.risk_score)
+        self.assertEqual(2, Vulnerability.objects.scope(self.dataspace).count())
+        self.assertEqual("8.4", str(package1.risk_score))
         self.assertEqual("8.4", str(product1.productpackages.get().weighted_risk_score))
 
+    def test_vulnerability_mixin_update_risk_score(self):
+        package1 = make_package(self.dataspace)
+
+        # Test with no vulnerabilities
+        package1.update_risk_score()
+        self.assertIsNone(package1.risk_score)
+
+        # Test with one vulnerability with risk score
+        vulnerability1 = make_vulnerability(dataspace=self.dataspace, risk_score=7.5)
+        vulnerability1.add_affected(package1)
+        package1.update_risk_score()
+        self.assertEqual("7.5", str(package1.risk_score))
+
+        # Test with multiple vulnerabilities, should use max
+        vulnerability2 = make_vulnerability(dataspace=self.dataspace, risk_score=9.2)
+        vulnerability2.add_affected(package1)
+        package1.update_risk_score()
+        self.assertEqual("9.2", str(package1.risk_score))
+
+        # Test with vulnerability with lower risk score, should keep max
+        vulnerability3 = make_vulnerability(dataspace=self.dataspace, risk_score=3.1)
+        vulnerability3.add_affected(package1)
+        package1.update_risk_score()
+        self.assertEqual("9.2", str(package1.risk_score))
+
+        # Test with all vulnerabilities having NULL risk scores
+        package2 = make_package(self.dataspace)
+        vulnerability4 = make_vulnerability(dataspace=self.dataspace, risk_score=None)
+        vulnerability5 = make_vulnerability(dataspace=self.dataspace, risk_score=None)
+        vulnerability4.add_affected(package2)
+        vulnerability5.add_affected(package2)
+        package2.update_risk_score()
+        self.assertIsNone(package2.risk_score)
+
     def test_vulnerability_model_affected_packages_m2m(self):
         package1 = make_package(self.dataspace)
         vulnerability1 = make_vulnerability(dataspace=self.dataspace, affecting=package1)