Add tests for pysec importer

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/tests/pipelines/test_github_importer_v2.py

@@ -1,36 +1,39 @@
-import pytest
-from unittest.mock import patch, MagicMock
 from datetime import datetime
+from unittest.mock import MagicMock
+from unittest.mock import patch
 
-from vulnerabilities.importer import AdvisoryData, AffectedPackage, Reference, VulnerabilitySeverity
-from vulnerabilities.utils import get_item
+import pytest
 from packageurl import PackageURL
 from univers.version_constraint import VersionConstraint
 from univers.version_range import RANGE_CLASS_BY_SCHEMES
 from univers.versions import SemverVersion
+
+from vulnerabilities.importer import AdvisoryData
+from vulnerabilities.importer import AffectedPackage
+from vulnerabilities.importer import Reference
+from vulnerabilities.importer import VulnerabilitySeverity
 from vulnerabilities.pipelines.v2_importers.github_importer import GitHubAPIImporterPipeline
-from vulnerabilities.pipelines.v2_importers.github_importer import get_cwes_from_github_advisory, get_purl
+from vulnerabilities.pipelines.v2_importers.github_importer import get_cwes_from_github_advisory
+from vulnerabilities.pipelines.v2_importers.github_importer import get_purl
+from vulnerabilities.utils import get_item
+
 
 @pytest.fixture
 def mock_fetch():
-    with patch("vulnerabilities.pipelines.v2_importers.github_importer.utils.fetch_github_graphql_query") as mock:
+    with patch(
+        "vulnerabilities.pipelines.v2_importers.github_importer.utils.fetch_github_graphql_query"
+    ) as mock:
         yield mock
 
 
 def test_advisories_count(mock_fetch):
     # Mock the GraphQL query response for advisory count
-    mock_fetch.return_value = {
-        "data": {
-            "securityVulnerabilities": {
-                "totalCount": 10
-            }
-        }
-    }
-    
+    mock_fetch.return_value = {"data": {"securityVulnerabilities": {"totalCount": 10}}}
+
     pipeline = GitHubAPIImporterPipeline()
-    
+
     count = pipeline.advisories_count()
-    
+
     # Assert that the count is correct
     assert count == 10
 
@@ -44,71 +47,66 @@ def test_collect_advisories(mock_fetch):
                     {
                         "node": {
                             "advisory": {
-                                "identifiers": [
-                                    {"type": "GHSA", "value": "GHSA-1234-ABCD"}
-                                ],
+                                "identifiers": [{"type": "GHSA", "value": "GHSA-1234-ABCD"}],
                                 "summary": "Sample advisory description",
-                                "references": [{"url": "https://github.com/advisories/GHSA-1234-ABCD"}],
+                                "references": [
+                                    {"url": "https://github.com/advisories/GHSA-1234-ABCD"}
+                                ],
                                 "severity": "HIGH",
-                                "cwes": {
-                                    "nodes": [{"cweId": "CWE-123"}]
-                                },
-                                "publishedAt": "2023-01-01T00:00:00Z"
+                                "cwes": {"nodes": [{"cweId": "CWE-123"}]},
+                                "publishedAt": "2023-01-01T00:00:00Z",
                             },
                             "firstPatchedVersion": {"identifier": "1.2.3"},
                             "package": {"name": "example-package"},
-                            "vulnerableVersionRange": ">=1.0.0,<=1.2.0"
+                            "vulnerableVersionRange": ">=1.0.0,<=1.2.0",
                         }
                     }
                 ],
-                "pageInfo": {
-                    "hasNextPage": False,
-                    "endCursor": None
-                }
+                "pageInfo": {"hasNextPage": False, "endCursor": None},
             }
         }
     }
 
     # Mock the response from GitHub GraphQL query
     mock_fetch.return_value = advisory_data
-    
+
     # Instantiate the pipeline
     pipeline = GitHubAPIImporterPipeline()
-    
+
     # Collect advisories
     advisories = list(pipeline.collect_advisories())
-    
+
     # Check if advisories were correctly parsed
     assert len(advisories) == 1
     advisory = advisories[0]
-    
+
     # Validate advisory fields
     assert advisory.advisory_id == "GHSA-1234-ABCD"
     assert advisory.summary == "Sample advisory description"
     assert advisory.url == "https://github.com/advisories/GHSA-1234-ABCD"
     assert len(advisory.references_v2) == 1
     assert advisory.references_v2[0].reference_id == "GHSA-1234-ABCD"
     assert advisory.severities[0].value == "HIGH"
-    
+
     # Validate affected package and version range
     affected_package = advisory.affected_packages[0]
     assert isinstance(affected_package.package, PackageURL)
     assert affected_package.package.name == "example-package"
-    
+
     # Check CWE extraction
     assert advisory.weaknesses == [123]
 
 
 def test_get_purl(mock_fetch):
     # Test for package URL generation
     result = get_purl("cargo", "example/package-name")
-    
+
     # Validate that the correct PackageURL is generated
     assert isinstance(result, PackageURL)
     assert result.type == "cargo"
     assert result.namespace == None
     assert result.name == "example/package-name"
-    
+
 
 def test_process_response(mock_fetch):
     # Mock advisory data as input for the process_response function
@@ -119,67 +117,58 @@ def test_process_response(mock_fetch):
                     {
                         "node": {
                             "advisory": {
-                                "identifiers": [
-                                    {"type": "GHSA", "value": "GHSA-5678-EFGH"}
-                                ],
+                                "identifiers": [{"type": "GHSA", "value": "GHSA-5678-EFGH"}],
                                 "summary": "Another advisory",
-                                "references": [{"url": "https://github.com/advisories/GHSA-5678-EFGH"}],
+                                "references": [
+                                    {"url": "https://github.com/advisories/GHSA-5678-EFGH"}
+                                ],
                                 "severity": "MEDIUM",
-                                "cwes": {
-                                    "nodes": [{"cweId": "CWE-200"}]
-                                },
-                                "publishedAt": "2023-02-01T00:00:00Z"
+                                "cwes": {"nodes": [{"cweId": "CWE-200"}]},
+                                "publishedAt": "2023-02-01T00:00:00Z",
                             },
                             "firstPatchedVersion": {"identifier": "2.0.0"},
                             "package": {"name": "another-package"},
-                            "vulnerableVersionRange": ">=2.0.0,<=3.0.0"
+                            "vulnerableVersionRange": ">=2.0.0,<=3.0.0",
                         }
                     }
                 ],
-                "pageInfo": {
-                    "hasNextPage": False,
-                    "endCursor": None
-                }
+                "pageInfo": {"hasNextPage": False, "endCursor": None},
             }
         }
     }
-    
+
     # Mock the response from GitHub GraphQL query
     mock_fetch.return_value = advisory_data
-    
+
     # Process the mock response
     result = list(GitHubAPIImporterPipeline().collect_advisories())
-    
+
     # Check the results
     assert len(result) == 1
     advisory = result[0]
-    
+
     # Validate the advisory data
     assert advisory.advisory_id == "GHSA-5678-EFGH"
     assert advisory.summary == "Another advisory"
     assert advisory.url == "https://github.com/advisories/GHSA-5678-EFGH"
-    
+
     # Check CWE extraction
     assert advisory.weaknesses == [200]
 
 
 def test_get_cwes_from_github_advisory(mock_fetch):
     # Mock CWEs extraction from GitHub advisory
-    advisory_data = {
-        "cwes": {
-            "nodes": [{"cweId": "CWE-522"}]
-        }
-    }
+    advisory_data = {"cwes": {"nodes": [{"cweId": "CWE-522"}]}}
 
     cwes = get_cwes_from_github_advisory(advisory_data)
-    
+
     # Validate the CWE ID extraction
     assert cwes == [522]
 
 
 def test_invalid_package_type_in_get_purl(mock_fetch):
     # Test for invalid package type
     result = get_purl("invalidpkg", "example/package-name")
-    
+
     # Assert that None is returned for an invalid package type
     assert result is None

vulnerabilities/tests/pipelines/test_pysec_v2_importer.py

@@ -0,0 +1,141 @@
+import json
+from io import BytesIO
+from unittest.mock import MagicMock
+from unittest.mock import patch
+from zipfile import BadZipFile
+from zipfile import ZipFile
+
+import pytest
+
+from vulnerabilities.importer import AdvisoryData
+from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipeline
+from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipelineV2
+from vulnerabilities.pipelines.v2_importers.pysec_importer import (
+    PyPIImporterPipeline,  # Path to the PyPI Importer
+)
+
+
+@pytest.fixture
+def mock_zip_data():
+    # Create mock zip data for testing
+    zip_buffer = BytesIO()
+    with ZipFile(zip_buffer, mode="w") as zip_file:
+        # Create a sample advisory file inside the zip
+        advisory_data = {
+            "advisory_id": "PYSEC-1234",
+            "summary": "Sample PyPI advisory",
+            "references": [{"url": "https://pypi.org/advisory/PYSEC-1234"}],
+            "package": {"name": "example-package"},
+            "affected_versions": ">=1.0.0,<=2.0.0",
+        }
+        # Save the sample advisory as a JSON file
+        with zip_file.open("PYSEC-1234.json", "w") as f:
+            f.write(json.dumps(advisory_data).encode("utf-8"))
+    zip_buffer.seek(0)
+    return zip_buffer
+
+
+@pytest.fixture
+def mock_requests_get():
+    with patch("requests.get") as mock:
+        yield mock
+
+
+def test_fetch_zip(mock_requests_get, mock_zip_data):
+    # Mock the `requests.get` to return the mock zip data
+    mock_requests_get.return_value.content = mock_zip_data.read()
+
+    pipeline = PyPIImporterPipeline()
+
+    # Call the `fetch_zip` method
+    pipeline.fetch_zip()
+
+    # Reset the position of mock_zip_data to 0 before comparing
+    mock_zip_data.seek(0)
+
+    # Verify that the zip file content is correctly assigned
+    assert pipeline.advisory_zip == mock_zip_data.read()
+
+
+def test_advisories_count(mock_requests_get, mock_zip_data):
+    # Mock the `requests.get` to return the mock zip data
+    mock_requests_get.return_value.content = mock_zip_data.read()
+
+    pipeline = PyPIImporterPipeline()
+
+    # Fetch the zip data
+    pipeline.fetch_zip()
+
+    # Test advisories count
+    count = pipeline.advisories_count()
+
+    # Verify that it correctly counts the number of advisory files starting with 'PYSEC-'
+    assert count == 1
+
+
+def test_collect_advisories(mock_requests_get, mock_zip_data):
+    # Mock the `requests.get` to return the mock zip data
+    mock_requests_get.return_value.content = mock_zip_data.read()
+
+    pipeline = PyPIImporterPipeline()
+
+    # Fetch the zip data
+    pipeline.fetch_zip()
+
+    # Mock the `parse_advisory_data_v2` function to return a dummy AdvisoryData
+    with patch("vulnerabilities.importers.osv.parse_advisory_data_v2") as mock_parse:
+        mock_parse.return_value = AdvisoryData(
+            advisory_id="PYSEC-1234",
+            summary="Sample PyPI advisory",
+            references_v2=[{"url": "https://pypi.org/advisory/PYSEC-1234"}],
+            affected_packages=[],
+            weaknesses=[],
+            url="https://pypi.org/advisory/PYSEC-1234",
+        )
+
+        # Call the `collect_advisories` method
+        advisories = list(pipeline.collect_advisories())
+
+        # Ensure we have 1 advisory
+        assert len(advisories) == 1
+
+        # Verify advisory data
+        advisory = advisories[0]
+        assert advisory.advisory_id == "PYSEC-1234"
+        assert advisory.summary == "Sample PyPI advisory"
+        assert advisory.url == "https://pypi.org/advisory/PYSEC-1234"
+
+
+def test_collect_advisories_invalid_file(mock_requests_get, mock_zip_data):
+    # Create a mock zip with an invalid file name
+    zip_buffer = BytesIO()
+    with ZipFile(zip_buffer, mode="w") as zip_file:
+        zip_file.writestr("INVALID_FILE.txt", "Invalid content")
+
+    zip_buffer.seek(0)
+    mock_requests_get.return_value.content = zip_buffer.read()
+
+    pipeline = PyPIImporterPipeline()
+
+    # Fetch the zip data
+    pipeline.fetch_zip()
+
+    # Mock the `parse_advisory_data_v2` function
+    with patch("vulnerabilities.importers.osv.parse_advisory_data_v2") as mock_parse:
+        mock_parse.return_value = AdvisoryData(
+            advisory_id="PYSEC-1234",
+            summary="Sample PyPI advisory",
+            references_v2=[{"url": "https://pypi.org/advisory/PYSEC-1234"}],
+            affected_packages=[],
+            weaknesses=[],
+            url="https://pypi.org/advisory/PYSEC-1234",
+        )
+
+        # Call the `collect_advisories` method and check the logging for invalid file
+        with patch(
+            "vulnerabilities.pipelines.VulnerableCodeBaseImporterPipelineV2.log"
+        ) as mock_log:
+            advisories = list(pipeline.collect_advisories())
+
+            # Ensure no advisories were yielded due to the invalid file
+            assert len(advisories) == 0