Fix tests

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/pipelines/v2_improvers/compute_package_risk.py

@@ -11,7 +11,7 @@
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.models import PackageV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
-from vulnerabilities.risk import compute_package_risk
+from vulnerabilities.risk import compute_package_risk_v2
 from vulnerabilities.risk import compute_vulnerability_risk_factors
 
 

vulnerabilities/risk.py

@@ -101,6 +101,22 @@ def compute_vulnerability_risk_factors(references, severities, exploits):
 
 
 def compute_package_risk(package):
+    """
+    Calculate the risk for a package by iterating over all vulnerabilities that affects this package
+    and determining the associated risk.
+    """
+    result = []
+    for relation in package.affectedbypackagerelatedvulnerability_set.all():
+        if risk := relation.vulnerability.risk_score:
+            result.append(float(risk))
+
+    if not result:
+        return
+
+    return round(max(result), 1)
+
+
+def compute_package_risk_v2(package):
     """
     Calculate the risk for a package by iterating over all vulnerabilities that affects this package
     and determining the associated risk.
@@ -113,4 +129,4 @@ def compute_package_risk(package):
     if not result:
         return
 
-    return round(max(result), 1)
+    return round(max(result), 1)