updates tests

Author: lauraschauer

prospector/core/prospector_test.py

@@ -2,19 +2,50 @@
 
 import pytest
 
+from llm.llm_service import LLMService
+
 from .prospector import prospector
 
 OPENCAST_CVE = "CVE-2021-21318"
 OPENCAST_REPO = "https://github.com/opencast/opencast"
 
 
+# Mock the llm_service configuration object
+class Config:
+    type: str = None
+    model_name: str = None
+    temperature: str = None
+    ai_core_sk: str = None
+
+    def __init__(self, type, model_name, temperature, ai_core_sk):
+        self.type = type
+        self.model_name = model_name
+        self.temperature = temperature
+        self.ai_core_sk = ai_core_sk
+
+
+config = Config("sap", "gpt-4", 0.0, "sk.json")
+
+
 def test_prospector_client():
     results, _ = prospector(
         vulnerability_id=OPENCAST_CVE,
         repository_url=OPENCAST_REPO,
         version_interval="9.1:9.2",
-        fetch_references=False,
         git_cache="/tmp/gitcache",
         limit_candidates=5000,
     )
     assert results[0].commit_id == "b18c6a7f81f08ed14884592a6c14c9ab611ad450"
+
+
+def test_prospector_llm_repo_url():
+    LLMService(config)
+
+    results, _ = prospector(
+        vulnerability_id=OPENCAST_CVE,
+        version_interval="9.1:9.2",
+        git_cache="/tmp/gitcache",
+        limit_candidates=5000,
+        use_llm_repository_url=True,
+    )
+    assert results[0].commit_id == "b18c6a7f81f08ed14884592a6c14c9ab611ad450"

prospector/data_sources/nvd/nvd_test.py

@@ -1,5 +1,5 @@
-from filter_entries import process_entries, retrieve_vulns
-from job_creation import enqueue_jobs
+from data_sources.nvd.filter_entries import process_entries, retrieve_vulns
+from data_sources.nvd.job_creation import enqueue_jobs
 
 # request new cves entries through NVD API and save to db
 cves = retrieve_vulns(7)

prospector/data_sources/nvd/version_extraction_test.py

@@ -19,9 +19,18 @@
             "nodes": [
                 {
                     "cpeMatch": [
-                        {"versionStartIncluding": "1.0", "versionEndIncluding": "2.0"},
-                        {"versionStartExcluding": "2.0", "versionEndExcluding": "3.0"},
-                        {"versionStartIncluding": "4.0", "versionEndIncluding": "5.0"},
+                        {
+                            "versionStartIncluding": "1.0",
+                            "versionEndIncluding": "2.0",
+                        },
+                        {
+                            "versionStartExcluding": "2.0",
+                            "versionEndExcluding": "3.0",
+                        },
+                        {
+                            "versionStartIncluding": "4.0",
+                            "versionEndIncluding": "5.0",
+                        },
                     ]
                 },
             ]
@@ -119,7 +128,7 @@ def test_extract_version_ranges_description():
     assert version_range == "None:8.0.4"
 
     version_range = extract_version_ranges_description(ADVISORY_TEXT_6)
-    assert version_range == "6.1.2:None"
+    assert version_range == "6.1.2.1:None"
 
 
 def test_extract_version_ranges_cpe():
@@ -145,4 +154,4 @@ def test_process_ranges():
 
 def test_extract_version_ranges():
     version_range = extract_version_range(JSON_DATA_4, ADVISORY_TEXT_6)
-    assert version_range == "6.1.2:None"
+    assert version_range == "6.1.2.1:None"