Feature/issue 110 (#119)

* Add some initial unit tests for stat and search

* Additional tests

Author: timmyteo

test/test_source.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from vdb.lib import VulnerabilityLocation
+from vdb.lib import VulnerabilityLocation, db6, search
 from vdb.lib.aqua import AquaSource
 from vdb.lib.cve import CVESource
 from vdb.lib.gha import GitHubSource
@@ -389,67 +389,166 @@ def test_aqua_wolfi_json():
     )
     with open(test_cve_data, "r") as fp:
         return json.loads(fp.read())
-
+    
 
 def test_convert(test_cve_json):
     nvdlatest = NvdSource()
-    data = nvdlatest.convert(test_cve_json)
-    assert len(data) == 384
-    for v in data:
+    vulnerabilities = nvdlatest.convert(test_cve_json)
+    assert len(vulnerabilities) == 384
+    for v in vulnerabilities:
         details = v.details
         for detail in details:
             assert detail
             assert detail.severity
             assert detail.package
             assert detail.package_type
+    
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 496
+    assert cve_index_count == 1155
+    results_count = len(list(search_db("CVE-2020-0001")))
+    assert results_count == 4
+    results_count = len(list(search_db("cpe:2.3:o:google:android:8.1:*:*:*:*:*:*:*")))
+    assert results_count == 25
+
     cvesource = CVESource()
-    data = cvesource.convert5(data)
-    assert len(data) == 384
+    cve = cvesource.convert5(vulnerabilities)
+    assert len(cve) == 384
+
+    db6.clear_all()
+    cvesource.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 0
+    assert cve_index_count == 0
 
 
 def test_convert2(test_cve_wconfig_json):
     nvdlatest = NvdSource()
-    data = nvdlatest.convert(test_cve_wconfig_json)
-    assert len(data) == 1
-    for v in data:
+    vulnerabilities = nvdlatest.convert(test_cve_wconfig_json)
+    assert len(vulnerabilities) == 1
+    for v in vulnerabilities:
         details = v.details
         for detail in details:
             assert detail
             assert detail.severity
             assert detail.package
             assert detail.package_type
             assert not detail.fixed_location
+
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 2
+    assert cve_index_count == 4
+    results_count = len(list(search_db("CVE-2020-8022")))
+    assert results_count == 4
+    results_count = len(list(search_db("cpe:2.3:o:opensuse:leap:15.1:*:*:*:*:*:*:*")))
+    assert results_count == 1
+
     cvesource = CVESource()
-    data = cvesource.convert5(data)
-    assert len(data) == 1
+    cve = cvesource.convert5(vulnerabilities)
+    assert len(cve) == 1
+
+    db6.clear_all()
+    cvesource.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 0
+    assert cve_index_count == 0
 
 
 def test_nvd_api_convert(test_nvd_api_json1, test_nvd_api_json2, test_nvd_api_json3, test_nvd_api_json4, test_nvd_api_git_json):
+    #json1
     nvdlatest = NvdSource()
-    data = nvdlatest.convert(test_nvd_api_json1)
-    assert len(data) == 1
-    for v in data:
+    vulnerabilities = nvdlatest.convert(test_nvd_api_json1)
+    assert len(vulnerabilities) == 1
+    for v in vulnerabilities:
         details = v.details
         for detail in details:
             assert detail
             assert detail.severity
             assert detail.package
             assert detail.package_type
             assert not detail.fixed_location
-    data = nvdlatest.convert(test_nvd_api_json2)
-    assert len(data) == 1
+    
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 4
+    assert cve_index_count == 20
+    results_count = len(list(search_db("CVE-2020-8022")))
+    assert results_count == 0
+    results_count = len(list(search_db("CVE-2024-0057")))
+    assert results_count == 10
+    results_count = len(list(search_db("cpe:2.3:a:microsoft:.net:*:*:*:*:*:*:*:*")))
+    assert results_count == 1
+
+    #json2
+    vulnerabilities = nvdlatest.convert(test_nvd_api_json2)
+    assert len(vulnerabilities) == 1
     cvesource = CVESource()
-    data = cvesource.convert5(data)
-    assert len(data) == 1
-    data = nvdlatest.convert(test_nvd_api_json3)
-    assert len(data) == 0
-    data = cvesource.convert5(data)
-    assert len(data) == 0
-    data = nvdlatest.convert(test_nvd_api_json4)
-    assert len(data) == 1
-    data = nvdlatest.convert(test_nvd_api_git_json)
-    assert len(data) == 1
-    assert len(data[0].details) == 2
+    cve = cvesource.convert5(vulnerabilities)
+    assert len(cve) == 1
+
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 1
+    assert cve_index_count == 7
+    results_count = len(list(search_db("CVE-2020-8022")))
+    assert results_count == 0
+    results_count = len(list(search_db("CVE-2024-21312")))
+    assert results_count == 7
+
+    #json3
+    vulnerabilities = nvdlatest.convert(test_nvd_api_json3)
+    assert len(vulnerabilities) == 0
+    cve = cvesource.convert5(vulnerabilities)
+    assert len(cve) == 0
+
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 0
+    assert cve_index_count == 0
+    results_count = len(list(search_db("CVE-2020-8022")))
+    assert results_count == 0
+    results_count = len(list(search_db("CVE-2024-23771")))
+    assert results_count == 0
+
+    #json4
+    vulnerabilities = nvdlatest.convert(test_nvd_api_json4)
+    assert len(vulnerabilities) == 1
+
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 2
+    assert cve_index_count == 21
+    results_count = len(list(search_db("CVE-2020-8022")))
+    assert results_count == 0
+    results_count = len(list(search_db("CVE-2015-3192")))
+    assert results_count == 21
+    results_count = len(list(search_db("cpe:2.3:a:pivotal_software:spring_framework:3.2.0:*:*:*:*:*:*:*")))
+    assert results_count == 2
+
+    #git_json
+    vulnerabilities = nvdlatest.convert(test_nvd_api_git_json)
+    assert len(vulnerabilities) == 1
+    assert len(vulnerabilities[0].details) == 2
+
+    db6.clear_all()
+    nvdlatest.store(vulnerabilities)
+    cve_data_count, cve_index_count = db6.stats()
+    assert cve_data_count == 2
+    assert cve_index_count == 2
+    results_count = len(list(search_db("CVE-2020-8022")))
+    assert results_count == 0
+    results_count = len(list(search_db("CVE-2023-52426")))
+    assert results_count == 2
+    results_count = len(list(search_db("cpe:2.3:a:libexpat_project:libexpat:*:*:*:*:*:*:*:*")))
+    assert results_count == 1
 
 
 @pytest.mark.skip(reason="This downloads and tests with live data")
@@ -654,3 +753,16 @@ def test_vuln_location():
     vl = VulnerabilityLocation.from_values("cpe:2.3:a:org.springframework:spring-web:*:*:*:*:*:*:*:*", "5.0.0.RC2",
                                            "*", "", "5.0.0.RC3")
     assert vl.version == ">=5.0.0.RC2-<5.0.0.RC3"
+
+
+def search_db(query):
+    if query.startswith("pkg:"):
+        results = search.search_by_purl_like(query, with_data=True)
+    elif query.startswith("CVE-") or query.startswith("GHSA-") or query.startswith("MAL-"):
+        results = search.search_by_cve(query, with_data=True)
+    elif query.startswith("http"):
+        results = search.search_by_url(query, with_data=True)
+    else:
+        results = search.search_by_cpe_like(query, with_data=True)
+    
+    return results