fix:补充了测试代码
This commit is contained in:
35
tests/test_CN_GPT_detection.py
Normal file
35
tests/test_CN_GPT_detection.py
Normal 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()
|
||||
Reference in New Issue
Block a user