Add a test for write_vul_data

Signed-off-by: ziadhany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/tests/test_export.py

@@ -1,15 +1,22 @@
 import os
 from io import StringIO
+from pathlib import Path
 from unittest import TestCase
 
 import pytest
+import saneyaml
 from django.core.management import call_command
 from django.core.management.base import CommandError
 
 from vulnerabilities.management.commands.export import create_sub_paths
+from vulnerabilities.models import Alias
 from vulnerabilities.models import Package
 from vulnerabilities.models import PackageRelatedVulnerability
 from vulnerabilities.models import Vulnerability
+from vulnerabilities.models import VulnerabilityReference
+from vulnerabilities.models import VulnerabilityRelatedReference
+from vulnerabilities.models import VulnerabilitySeverity
+from vulnerabilities.models import Weakness
 
 
 @pytest.mark.parametrize(
@@ -41,11 +48,38 @@ def package(db):
 
 
 @pytest.fixture
-def vulnerability(db):
-    return Vulnerability.objects.create(
+def vulnerability_reference():
+    return VulnerabilityReference.objects.create(
+        reference_id="fake",
+        url=f"https://..",
+    )
+
+
+@pytest.fixture
+def vulnerability_severity(vulnerability_reference):
+    return VulnerabilitySeverity.objects.create(
+        scoring_system="cvssv3_vector",
+        value="CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H",
+        reference_id=vulnerability_reference.id,
+    )
+
+
+@pytest.fixture
+def vulnerability(db, vulnerability_reference, vulnerability_severity):
+    vulnerability = Vulnerability.objects.create(
         vulnerability_id="VCID-pst6-b358-aaap",
         summary="test-vuln",
     )
+    Alias.objects.create(alias=f"CVE-xxx-xxx-xx", vulnerability=vulnerability)
+
+    VulnerabilityRelatedReference.objects.create(
+        reference=vulnerability_reference, vulnerability=vulnerability
+    )
+
+    weakness = Weakness.objects.create(cwe_id=15)
+    vulnerability.weaknesses.add(weakness)
+
+    return vulnerability
 
 
 @pytest.fixture
@@ -67,9 +101,50 @@ def test_missing_path(self):
         assert "Error: the following arguments are required: path" in err
 
     def test_bad_path_fail_error(self):
-        buf = StringIO()
         with pytest.raises(CommandError) as cm:
-            call_command("export", "/bad path", stdout=buf)
+            call_command("export", "/bad path", stdout=StringIO())
 
         err = str(cm)
         assert "Please enter a valid path" in err
+
+
+@pytest.mark.django_db
+def test_write_vul_data(
+    tmp_path, package_related_vulnerability, vulnerability_reference, vulnerability_severity
+):
+    expected_vul = {
+        "vulnerability_id": "VCID-pst6-b358-aaap",
+        "aliases": ["CVE-xxx-xxx-xx"],
+        "summary": "test-vuln",
+        "severities": [
+            {
+                "id": vulnerability_severity.id,
+                "reference_id": vulnerability_reference.id,
+                "scoring_system": "cvssv3_vector",
+                "value": "CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H",
+                "scoring_elements": "",
+            }
+        ],
+        "references": [
+            {"id": vulnerability_reference.id, "url": "https://..", "reference_id": "fake"}
+        ],
+        "weaknesses": ["CWE-15"],
+    }
+    expected_pkg = {
+        "package": "pkg:generic/nginx/test",
+        "versions": [
+            {
+                "purl": "pkg:generic/nginx/test@2",
+                "affected_by_vulnerabilities": ["VCID-pst6-b358-aaap"],
+                "fixing_vulnerabilities": [],
+            },
+        ],
+    }
+
+    call_command("export", tmp_path, stdout=StringIO())
+    # path: type/namespace/name
+    vul_filepath = os.path.join(tmp_path, "generic/nginx/test/VCID-pst6-b358-aaap.yml")
+    pck_filepath = os.path.join(tmp_path, "generic/nginx/test/generic-nginx-test.yml")
+
+    assert Path(vul_filepath).read_text() == saneyaml.dump(expected_vul)
+    assert Path(pck_filepath).read_text() == saneyaml.dump(expected_pkg)