Refactor tests

Signed-off-by: Shivam Sandbhor <shivam.sandbhor@gmail.com>

Author: sbs2001

vulnerabilities/importers/elixir_security.py

@@ -87,10 +87,10 @@ def get_versions_for_pkg_from_range_list(self, version_range_list, pkg_name):
         all_version_list = self.pkg_manager_api.get(pkg_name)
         if not version_range_list:
             return [], all_version_list
-        version_ranges = {
+        version_ranges = [
             VersionSpecifier.from_scheme_version_spec_string("semver", r)
             for r in version_range_list
-        }
+        ]
         for version in all_version_list:
             version_obj = SemverVersion(version)
             if any([version_obj in v for v in version_ranges]):

vulnerabilities/importers/rust.py

@@ -23,7 +23,9 @@
 import asyncio
 import re
 from itertools import chain
-from typing import Optional, Mapping
+from typing import Optional
+from typing import Mapping
+from typing import List
 from typing import Set
 from typing import Tuple
 from urllib.error import HTTPError
@@ -104,22 +106,22 @@ def _load_advisory(self, path: str) -> Optional[Advisory]:
 
         # FIXME: Avoid wildcard version ranges for now.
         # See https://github.com/RustSec/advisory-db/discussions/831
-        affected_ranges = {
+        affected_ranges = [
             VersionSpecifier.from_scheme_version_spec_string("semver", r)
             for r in chain.from_iterable(record.get("affected", {}).get("functions", {}).values())
             if r != "*"
-        }
+        ]
 
-        unaffected_ranges = {
+        unaffected_ranges = [
             VersionSpecifier.from_scheme_version_spec_string("semver", r)
             for r in record.get("versions", {}).get("unaffected", [])
             if r != "*"
-        }
-        resolved_ranges = {
+        ]
+        resolved_ranges = [
             VersionSpecifier.from_scheme_version_spec_string("semver", r)
             for r in record.get("versions", {}).get("patched", [])
             if r != "*"
-        }
+        ]
 
         unaffected, affected = categorize_versions(
             all_versions, unaffected_ranges, affected_ranges, resolved_ranges
@@ -153,9 +155,9 @@ def _load_advisory(self, path: str) -> Optional[Advisory]:
 
 def categorize_versions(
     all_versions: Set[str],
-    unaffected_versions: Set[VersionSpecifier],
-    affected_versions: Set[VersionSpecifier],
-    resolved_versions: Set[VersionSpecifier],
+    unaffected_version_ranges: List[VersionSpecifier],
+    affected_version_ranges: List[VersionSpecifier],
+    resolved_version_ranges: List[VersionSpecifier],
 ) -> Tuple[Set[str], Set[str]]:
     """
     Categorize all versions of a crate according to the given version ranges.
@@ -165,24 +167,30 @@ def categorize_versions(
 
     unaffected, affected = set(), set()
 
-    if not any(unaffected_versions.union(affected_versions).union(resolved_versions)):
+    if (
+        not unaffected_version_ranges
+        and not affected_version_ranges
+        and not resolved_version_ranges
+    ):
         return unaffected, affected
 
     # TODO: This is probably wrong
     for version in all_versions:
         version_obj = SemverVersion(version)
-        if affected_versions and all([version_obj in av for av in affected_versions]):
+        if affected_version_ranges and all([version_obj in av for av in affected_version_ranges]):
             affected.add(version)
-        elif unaffected_versions and all([version_obj in av for av in unaffected_versions]):
+        elif unaffected_version_ranges and all(
+            [version_obj in av for av in unaffected_version_ranges]
+        ):
             unaffected.add(version)
-        elif resolved_versions and all([version_obj in av for av in resolved_versions]):
+        elif resolved_version_ranges and all([version_obj in av for av in resolved_version_ranges]):
             unaffected.add(version)
 
     # If some versions were not classified above, one or more of the given ranges might be empty, so
     # the remaining versions default to either affected or unaffected.
     uncategorized_versions = all_versions - unaffected.union(affected)
     if uncategorized_versions:
-        if not affected_versions:
+        if not affected_version_ranges:
             affected.update(uncategorized_versions)
         else:
             unaffected.update(uncategorized_versions)

vulnerabilities/importers/safety_db.py

@@ -167,16 +167,19 @@ def categorize_versions(
     :return: impacted, resolved purls
     """
     impacted_versions, impacted_purls = set(), set()
-    vurl_spec = f"pypi:{ ','.join(version_specs) }"
-    vurl_spec = VersionSpecifier.from_version_spec_string(vurl_spec)
+    # vurl_spec = f"pypi:{ ','.join(version_specs) }"
+    # vurl_spec = VersionSpecifier.from_version_spec_string(vurl_spec)
+    vurl_specs = []
+    for version_spec in version_specs:
+        vurl_specs.append(VersionSpecifier.from_scheme_version_spec_string("pypi", version_spec))
 
     for version in all_versions:
         try:
             version_object = PYPIVersion(version)
         except:
             continue
 
-        if version_object in vurl_spec:
+        if any([version_object in vurl_spec for vurl_spec in vurl_specs]):
             impacted_versions.add(version)
             impacted_purls.add(
                 PackageURL(

vulnerabilities/tests/test_rust.py

@@ -52,12 +52,12 @@
 def test_categorize_versions():
     flatbuffers_versions = MOCKED_CRATES_API_VERSIONS["flatbuffers"]
 
-    unaffected_ranges = {VersionSpecifier.from_scheme_version_spec_string("semver", "< 0.4.0")}
-    affected_ranges = {
+    unaffected_ranges = [VersionSpecifier.from_scheme_version_spec_string("semver", "< 0.4.0")]
+    affected_ranges = [
         VersionSpecifier.from_scheme_version_spec_string("semver", ">= 0.4.0"),
         VersionSpecifier.from_scheme_version_spec_string("semver", "<= 0.6.0"),
-    }
-    resolved_ranges = {VersionSpecifier.from_scheme_version_spec_string("semver", ">= 0.6.1")}
+    ]
+    resolved_ranges = [VersionSpecifier.from_scheme_version_spec_string("semver", ">= 0.6.1")]
 
     unaffected_versions, affected_versions = categorize_versions(
         set(flatbuffers_versions),
@@ -77,9 +77,9 @@ def test_categorize_versions():
 def test_categorize_versions_without_affected_ranges():
     all_versions = {"1.0", "1.1", "2.0", "2.1", "3.0", "3.1"}
 
-    unaffected_ranges = {VersionSpecifier.from_scheme_version_spec_string("semver", "< 1.2")}
-    affected_ranges = set()
-    resolved_ranges = {VersionSpecifier.from_scheme_version_spec_string("semver", ">= 3.0")}
+    unaffected_ranges = [VersionSpecifier.from_scheme_version_spec_string("semver", "< 1.2")]
+    affected_ranges = []
+    resolved_ranges = [VersionSpecifier.from_scheme_version_spec_string("semver", ">= 3.0")]
 
     unaffected_versions, affected_versions = categorize_versions(
         all_versions,
@@ -102,12 +102,12 @@ def test_categorize_versions_without_affected_ranges():
 def test_categorize_versions_with_only_affected_ranges():
     all_versions = {"1.0", "1.1", "2.0", "2.1", "3.0", "3.1"}
 
-    unaffected_ranges = set()
-    affected_ranges = {
+    unaffected_ranges = []
+    affected_ranges = [
         VersionSpecifier.from_scheme_version_spec_string("semver", "> 1.2"),
         VersionSpecifier.from_scheme_version_spec_string("semver", "<= 2.1"),
-    }
-    resolved_ranges = set()
+    ]
+    resolved_ranges = []
 
     unaffected_versions, affected_versions = categorize_versions(
         all_versions,
@@ -130,7 +130,7 @@ def test_categorize_versions_with_only_affected_ranges():
 def test_categorize_versions_without_any_ranges():
     all_versions = {"1.0", "1.1", "2.0", "2.1", "3.0", "3.1"}
 
-    unaffected, affected = categorize_versions(all_versions, set(), set(), set())
+    unaffected, affected = categorize_versions(all_versions, [], [], [])
 
     assert len(unaffected) == 0
     assert len(affected) == 0

vulnerabilities/tests/test_safety_db.py

@@ -74,10 +74,10 @@ def test_import(self):
             with patch("vulnerabilities.importers.SafetyDbDataSource.set_api"):
                 runner.run()
 
-        assert models.Vulnerability.objects.count() == 9
-        assert models.VulnerabilityReference.objects.count() == 9
-        assert models.PackageRelatedVulnerability.objects.filter(is_vulnerable=False).count() == 18
-        assert models.PackageRelatedVulnerability.objects.filter(is_vulnerable=True).count() == 18
+        assert models.Vulnerability.objects.count() == 8
+        assert models.VulnerabilityReference.objects.count() == 8
+        assert models.PackageRelatedVulnerability.objects.filter(is_vulnerable=False).count() == 15
+        assert models.PackageRelatedVulnerability.objects.filter(is_vulnerable=True).count() == 17
 
         expected_package_count = sum([len(v) for v in MOCK_VERSION_API.cache.values()])
         assert models.Package.objects.count() == expected_package_count
@@ -106,13 +106,6 @@ def test_import(self):
             cve_ids={"CVE-2016-6186"},
         )
 
-        self.assert_by_vulnerability(
-            "pyup.io-38115",
-            "zulip",
-            {"2.0"},
-            {"2.1.1", "2.1.2", "2.1.3"},
-        )
-
         self.assert_by_vulnerability(
             "pyup.io-38114",
             "zulip",