Add tests for vulnerability endpoint

And address review comments

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

Author: TG1999

vulnerabilities/api.py

@@ -9,6 +9,8 @@
 
 from urllib.parse import unquote
 
+from django.db.models import Prefetch
+from django.db.models import Q
 from django_filters import rest_framework as filters
 from packageurl import PackageURL
 from rest_framework import serializers
@@ -50,34 +52,7 @@ class Meta:
         fields = ["url", "purl"]
 
 
-class FilteredPackageListSerializer(serializers.ListSerializer):
-    def to_representation(self, data):
-        params = self.context["request"].query_params
-        name = params.get("name")
-        if name:
-            data = data.filter(name=name)
-        namespace = params.get("namespace")
-        if namespace:
-            data = data.filter(namespace=namespace)
-        type = params.get("type")
-        if type:
-            data = data.filter(type=type)
-        return super(FilteredPackageListSerializer, self).to_representation(data)
-
-
-class FixedPackageSerializer(serializers.ModelSerializer):
-
-    purl = serializers.CharField(source="package_url")
-
-    class Meta:
-        list_serializer_class = FilteredPackageListSerializer
-        model = Package
-        fields = ["url", "purl"]
-
-
-class MinimalVulnerabilitySerializerWithReferencesAndSummary(
-    serializers.HyperlinkedModelSerializer
-):
+class VulnSerializerRefsAndSummary(serializers.HyperlinkedModelSerializer):
     """
     Used for nesting inside package focused APIs.
     """
@@ -99,7 +74,7 @@ class Meta:
         fields = ["url", "vulnerability_id"]
 
 
-class MinimalPackageSerializerWithFixedVulnerabilities(serializers.HyperlinkedModelSerializer):
+class PackageSerializerFixedVulns(serializers.HyperlinkedModelSerializer):
     """
     Used for nesting inside vulnerability focused APIs.
     """
@@ -126,7 +101,9 @@ class Meta:
 
 class VulnerabilitySerializer(serializers.HyperlinkedModelSerializer):
 
-    fixed_packages = FixedPackageSerializer(many=True, source="resolved_to", read_only=True)
+    fixed_packages = MinimalPackageSerializer(
+        many=True, source="filtered_fixed_packages", read_only=True
+    )
     affected_packages = MinimalPackageSerializer(many=True, source="vulnerable_to", read_only=True)
 
     references = VulnerabilityReferenceSerializer(many=True, source="vulnerabilityreference_set")
@@ -152,13 +129,13 @@ def to_representation(self, instance):
         return data
 
     purl = serializers.CharField(source="package_url")
-    affected_by_vulnerabilities = MinimalVulnerabilitySerializerWithReferencesAndSummary(
+    affected_by_vulnerabilities = VulnSerializerRefsAndSummary(
         many=True, source="vulnerable_to", read_only=True
     )
-    fixing_vulnerabilities = MinimalVulnerabilitySerializerWithReferencesAndSummary(
+    fixing_vulnerabilities = VulnSerializerRefsAndSummary(
         many=True, source="resolved_to", read_only=True
     )
-    fixed_packages = MinimalPackageSerializerWithFixedVulnerabilities(many=True, read_only=True)
+    fixed_packages = PackageSerializerFixedVulns(many=True, read_only=True)
 
     class Meta:
         model = Package
@@ -247,7 +224,27 @@ class Meta:
 
 
 class VulnerabilityViewSet(viewsets.ReadOnlyModelViewSet):
-    queryset = Vulnerability.objects.all()
+    def get_queryset(self):
+        params = self.request.query_params
+        query = Q()
+        name = params.get("name")
+        if name:
+            query &= Q(name=name)
+        namespace = params.get("namespace")
+        if namespace:
+            query &= Q(namespace=namespace)
+        type = params.get("type")
+        if type:
+            query &= Q(type=type)
+        queryset = Vulnerability.objects.prefetch_related(
+            Prefetch(
+                "packages",
+                queryset=Package.objects.filter(query, packagerelatedvulnerability__fix=True),
+                to_attr="filtered_fixed_packages",
+            )
+        )
+        return queryset
+
     serializer_class = VulnerabilitySerializer
     paginate_by = 50
     filter_backends = (filters.DjangoFilterBackend,)

vulnerabilities/tests/test_fix_api.py

@@ -26,6 +26,12 @@ def setUp(self):
                 summary=str(i),
             )
         self.vulnerability = Vulnerability.objects.create(summary="test")
+        self.pkg1 = Package.objects.create(name="flask", type="pypi", version="0.1.2")
+        self.pkg2 = Package.objects.create(name="flask", type="debian", version="0.1.2")
+        for pkg in [self.pkg1, self.pkg2]:
+            PackageRelatedVulnerability.objects.create(
+                package=pkg, vulnerability=self.vulnerability, fix=True
+            )
 
     def test_api_status(self):
         response = self.client.get("/api/vulnerabilities/", format="json")
@@ -44,7 +50,35 @@ def test_api_with_single_vulnerability(self):
             "vulnerability_id": f"VULCOID-{int_to_base36(self.vulnerability.id).upper()}",
             "summary": "test",
             "aliases": [],
-            "fixed_packages": [],
+            "fixed_packages": [
+                {
+                    "url": f"http://testserver/api/packages/{self.pkg1.id}",
+                    "purl": "pkg:pypi/[email protected]",
+                },
+                {
+                    "url": f"http://testserver/api/packages/{self.pkg2.id}",
+                    "purl": "pkg:debian/[email protected]",
+                },
+            ],
+            "affected_packages": [],
+            "references": [],
+        }
+
+    def test_api_with_single_vulnerability_with_filters(self):
+        response = self.client.get(
+            f"/api/vulnerabilities/{self.vulnerability.id}?type=pypi", format="json"
+        ).data
+        assert response == {
+            "url": f"http://testserver/api/vulnerabilities/{self.vulnerability.id}",
+            "vulnerability_id": f"VULCOID-{int_to_base36(self.vulnerability.id).upper()}",
+            "summary": "test",
+            "aliases": [],
+            "fixed_packages": [
+                {
+                    "url": f"http://testserver/api/packages/{self.pkg1.id}",
+                    "purl": "pkg:pypi/[email protected]",
+                },
+            ],
             "affected_packages": [],
             "references": [],
         }