diff --git a/tests/test_final_tests.py b/tests/test_final_tests.py index e7bbcb1..98733ab 100644 --- a/tests/test_final_tests.py +++ b/tests/test_final_tests.py @@ -1,5 +1,7 @@ +import time import unittest import shutil +import threading from detection.utils import read_file_content from .final_tests_util import clone_repo, Path, inject_random_backdoor @@ -7,6 +9,25 @@ from detection.Regexdetection import find_dangerous_functions from detection.GPTdetection import detectGPT +def GPTdetectFileList(fileList): + results = [] + threads = [] + for file in fileList: + content = read_file_content(str(file)) + threads.append(threading.Thread(target=GPTThread(), args=(content, results))) + for thread in threads: + thread.start() + time.sleep(0.5) + for thread in threads: + thread.join() + return results + +def GPTThread(content, results): + try: + results.append(detectGPT(content)) + except Exception as e: + print(e) + class TestFinalTests(unittest.TestCase): def setUp(self) -> None: self.path = "./tmp/repo" @@ -39,23 +60,12 @@ class TestFinalTests(unittest.TestCase): print(detectedNum / self.py_filesNum) self.assertAlmostEqual(detectedNum, self.py_filesNum, places=1) GPTdetectedNum = 0 - - for i in possibly_dangerous_file: - content = read_file_content(str(i)) - results = {} - try: - results = detectGPT(content) - if ( - len(results["high"]) > 0 - or len(results["medium"]) > 0 - or len(results["low"]) > 0 - ): - GPTdetectedNum += 1 - print(GPTdetectedNum) - self.assertGreaterEqual(GPTdetectedNum, detectedNum) - - except Exception as e: - print(e) + GPTresult = GPTdetectFileList(possibly_dangerous_file) + for result in GPTresult: + if len(result) > 0: + GPTdetectedNum += 1 + print(GPTdetectedNum) + self.assertGreaterEqual(GPTdetectedNum, detectedNum) # test pickle files pickle_detectedNum = 0