fix:补充了测试代码

This commit is contained in:
2024-05-15 13:38:01 +08:00
parent 2e5460a522
commit 9d6f054478
2 changed files with 132 additions and 0 deletions

View File

@@ -0,0 +1,35 @@
import unittest
import warnings
import os
from detection.cngptdetection import cndetectGPT
class TestBackdoorDetection(unittest.TestCase):
def test_gpt_risk_detection(self):
content = """import os
os.system('rm -rf /') # high risk
exec('print("Hello")') # high risk
eval('2 + 2') # high risk
"""
results = cndetectGPT(content)
self.assertEqual(len(results["high"]), 3)
def test_gpt_no_risk_detection(self):
content = """a = 10
b = a + 5
print('This should not be detected as risky.')
"""
results = cndetectGPT(content)
self.assertEqual(len(results["high"]), 0)
self.assertEqual(len(results["medium"]), 0)
self.assertEqual(len(results["low"]), 0)
def test_gpt_env_no_set(self):
if os.getenv("BAIDU_API_KEY") is None or os.getenv("BAIDU_SECRET_KEY") is None:
self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
content = "print('test test')"
with self.assertRaises(ValueError):
cndetectGPT(content)
if __name__ == "__main__":
unittest.main()