Fix failing tests

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

Author: TG1999

vulnerabilities/models.py

@@ -317,28 +317,40 @@ class Meta:
         ordering = ["vulnerability", "reference"]
 
 
+def purl_to_dict(purl: PackageURL):
+    """
+    Return a dict of purl components suitable for use in a queryset.
+    We need to have specific empty values for using in querysets because of our peculiar model structure.
+
+    For example::
+    >>> purl_to_dict(PackageURL.from_string("pkg:generic/postgres"))
+    {'type': 'generic', 'namespace': '', 'name': 'postgres', 'version': '', 'qualifiers': {}, 'subpath': ''}
+    >>> purl_to_dict(PackageURL.from_string("pkg:generic/postgres/[email protected]?foo=bar#baz"))
+    {'type': 'generic', 'namespace': 'postgres', 'name': 'postgres', 'version': '1.2', 'qualifiers': {'foo': 'bar'}, 'subpath': 'baz'}
+    """
+    if isinstance(purl, str):
+        purl = PackageURL.from_string(purl)
+
+    return dict(
+        type=purl.type,
+        namespace=purl.namespace or "",
+        name=purl.name,
+        version=purl.version or "",
+        qualifiers=purl.qualifiers or {},
+        subpath=purl.subpath or "",
+    )
+
+
 class PackageQuerySet(BaseQuerySet, PackageURLQuerySet):
     def get_or_create_from_purl(self, purl: PackageURL):
         """
         Return an existing or new Package (created if neeed) given a
         ``purl`` PackageURL.
         """
-        purl_fields = without_empty_values(purl.to_dict(encode=True))
-
-        # when there are 2 packages one with qualifiers and one without
-        # qualifiers, having all other fields same, this raises MultipleObjectsReturned
-        # so we are filling out the fields with empty value to avoid this
-        # for field in PackageURL._fields:
-        #     # name, type, and version are required fields
-        #     if field not in purl_fields:
-        #         if field == "namespace":
-        #             purl_fields[field] = ""
-        #         if field == "qualifiers":
-        #             purl_fields[field] = {}
-        #         if field == "subpath":
-        #             purl_fields[field] = ""
-
-        package, _ = Package.objects.get_or_create(**purl_fields)
+        if isinstance(purl, str):
+            purl = PackageURL.from_string(purl)
+
+        package, _ = Package.objects.get_or_create(**purl_to_dict(purl=purl))
         return package
 
     def for_package_url_object(self, purl):
@@ -347,15 +359,12 @@ def for_package_url_object(self, purl):
         ``purl`` string is validated and transformed into filtering lookups. If
         this is a PackageURL object it is reused as-is.
         """
-        if isinstance(purl, PackageURL):
-            lookups = without_empty_values(purl.to_dict(encode=True))
-            return self.filter(**lookups)
-
-        elif isinstance(purl, str):
-            return self.for_package_url(purl, encode=False)
-
-        else:
+        if not purl:
             return self.none()
+        if isinstance(purl, str):
+            purl = PackageURL.from_string(purl)
+        lookups = without_empty_values(purl.to_dict(encode=True))
+        return self.filter(**lookups)
 
     def affected(self):
         """

vulnerabilities/tests/test_postgresql.py

@@ -7,23 +7,13 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 
-import json
 import os
 
-from django.test import TestCase
-from packageurl import PackageURL
+import pytest
 
-from vulnerabilities import severity_systems
-from vulnerabilities.import_runner import process_advisories
-from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importer import Reference
-from vulnerabilities.importer import VulnerabilitySeverity
 from vulnerabilities.importers.postgresql import to_advisories
-from vulnerabilities.improve_runner import ImproveRunner
-from vulnerabilities.improve_runner import process_inferences
-from vulnerabilities.improvers.default import DefaultImprover
+from vulnerabilities.models import Package
 from vulnerabilities.tests import util_tests
-from vulnerabilities.utils import AffectedPackage
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 TEST_DATA = os.path.join(
@@ -32,18 +22,19 @@
 )
 
 
-class TestPostgreSQLImporter(TestCase):
-    def test_to_advisories(self):
-        with open(os.path.join(TEST_DATA, "advisories.html")) as f:
-            raw_data = f.read()
-        advisories = to_advisories(raw_data)
-        result = [data.to_dict() for data in advisories]
-        expected_file = os.path.join(TEST_DATA, f"parse-advisory-postgresql-expected.json")
-        util_tests.check_results_against_json(result, expected_file)
-
-    def test_run_default_improver(self):
-        with open(os.path.join(TEST_DATA, "improver-data.json")) as f:
-            raw_data = json.load(f)
-        advisories = [AdvisoryData.from_dict(data) for data in raw_data]
-        process_advisories(advisories, "postgresql")
-        ImproveRunner(DefaultImprover).run()
+def test_to_advisories():
+    with open(os.path.join(TEST_DATA, "advisories.html")) as f:
+        raw_data = f.read()
+    advisories = to_advisories(raw_data)
+    result = [data.to_dict() for data in advisories]
+    expected_file = os.path.join(TEST_DATA, f"parse-advisory-postgresql-expected.json")
+    util_tests.check_results_against_json(result, expected_file)
+
+
+@pytest.mark.django_db
+def test_get_or_create_from_purl():
+    p1 = "pkg:generic/postgres"
+    p2 = "pkg:generic/postgres?foo=bar"
+    res1 = Package.objects.get_or_create_from_purl(p1)
+    res2 = Package.objects.get_or_create_from_purl(p2)
+    res3 = Package.objects.get_or_create_from_purl(p1)