34 lines
1.3 KiB
Python
Executable File
34 lines
1.3 KiB
Python
Executable File
import unittest
|
|
from src.preprocessor import Preprocessor
|
|
|
|
class TestPreprocessor(unittest.TestCase):
|
|
def setUp(self):
|
|
self.preprocessor = Preprocessor(max_length=100)
|
|
|
|
def test_prepare_data(self):
|
|
cve_data = [
|
|
("project1", "CVE-2021-1234", "CWE-79", "fix1", ["intro1"]),
|
|
("project2", "CVE-2021-5678", "CWE-89", "fix2", ["intro2", "intro3"])
|
|
]
|
|
|
|
prompts, labels = self.preprocessor.prepare_data(cve_data)
|
|
|
|
expected_prompts = [
|
|
"Project: project1\nCVE ID: CVE-2021-1234\nCWE: CWE-79\nFixing Commit: fix1\nIdentify the vulnerability introducing commit:",
|
|
"Project: project2\nCVE ID: CVE-2021-5678\nCWE: CWE-89\nFixing Commit: fix2\nIdentify the vulnerability introducing commit:"
|
|
]
|
|
expected_labels = [["intro1"], ["intro2", "intro3"]]
|
|
|
|
self.assertEqual(prompts, expected_prompts)
|
|
self.assertEqual(labels, expected_labels)
|
|
|
|
def test_max_length_truncation(self):
|
|
self.preprocessor.max_length = 50
|
|
cve_data = [("very_long_project_name", "CVE-2021-1234", "CWE-79", "very_long_fixing_commit_hash", ["intro1"])]
|
|
|
|
prompts, _ = self.preprocessor.prepare_data(cve_data)
|
|
|
|
self.assertEqual(len(prompts[0]), 50)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main() |