From fbeba5b4fc44268c2043d04246cb2ba7d677c516 Mon Sep 17 00:00:00 2001
From: sangge-redmi <2251250136@qq.com>
Date: Tue, 4 Jun 2024 15:05:18 +0800
Subject: [PATCH] feat: update test cases

---
 tests/test_final_tests.py | 31 ++++++++++++++++++++++---------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/tests/test_final_tests.py b/tests/test_final_tests.py
index 53075dd..e7bbcb1 100644
--- a/tests/test_final_tests.py
+++ b/tests/test_final_tests.py
@@ -9,19 +9,21 @@ from detection.GPTdetection import detectGPT
 
 class TestFinalTests(unittest.TestCase):
     def setUp(self) -> None:
-        shutil.rmtree("./tmp/repo", ignore_errors=True)
-        clone_repo("https://github.com/injetlee/Python.git", "./tmp/repo")
+        self.path = "./tmp/repo"
+        shutil.rmtree(self.path, ignore_errors=True)
+        clone_repo("https://github.com/injetlee/Python.git", self.path)
         sampleRate = 0.1
-        self.inject_reslt = inject_random_backdoor("./tmp/repo", sample_rate=sampleRate)
+        self.inject_reslt = inject_random_backdoor(self.path, sample_rate=sampleRate)
         self.injectedNum = len(self.inject_reslt)
         print(self.injectedNum)
-        project_path = Path("./tmp/repo")
+        project_path = Path(self.path)
         self.all_python_files = list(project_path.rglob("*.py"))
-        self.filesNum = len(self.all_python_files)
-        self.trueRate = self.injectedNum / self.filesNum
+        self.py_filesNum = len(self.all_python_files)
+        self.trueRate = self.injectedNum / self.py_filesNum
         print(self.trueRate)
 
-    def test_final_tests(self):
+    # test backdoor code in python files
+    def test_final_tests_pycode(self):
         detectedNum = 0
         possibly_dangerous_file = []
         for file in self.all_python_files:
@@ -34,8 +36,8 @@ class TestFinalTests(unittest.TestCase):
             ):
                 detectedNum += 1
                 possibly_dangerous_file.append(file)
-        print(detectedNum / self.filesNum)
-        self.assertAlmostEqual(detectedNum / self.filesNum, self.trueRate, places=1)
+        print(detectedNum / self.py_filesNum)
+        self.assertAlmostEqual(detectedNum, self.py_filesNum, places=1)
         GPTdetectedNum = 0
 
         for i in possibly_dangerous_file:
@@ -55,6 +57,17 @@ class TestFinalTests(unittest.TestCase):
             except Exception as e:
                 print(e)
 
+        # test pickle files
+        pickle_detectedNum = 0
+        pickle_tureNum = len(list(Path(self.path).glob("*.pickle")))
+
+        self.assertAlmostEqual(pickle_detectedNum, pickle_tureNum, places=1)
+
+        # test pyc files
+        pyc_detectedNum = 0
+        pyc_tureNum = len(list(Path(self.path).glob("*.pyc")))
+        self.assertAlmostEqual(pyc_detectedNum, pyc_tureNum, places=1)
+
 
 if __name__ == "__main__":
     unittest.main()