50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
import unittest
|
|
import warnings
|
|
import os
|
|
|
|
from detection.cngptdetection import detectGPT # 导入调用百度 ai 模型的函数
|
|
|
|
|
|
class TestBackdoorDetection(unittest.TestCase):
|
|
def test_gpt_risk_detection(self):
|
|
"""
|
|
if os.getenv("BAIDU_API_KEY") is None or os.getenv("BAIDU_SECRET_KEY") is None:
|
|
warnings.warn("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set, test skipped.", UserWarning)
|
|
self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
|
|
"""
|
|
content = """import os
|
|
os.system('rm -rf /') # high risk
|
|
exec('print("Hello")') # high risk
|
|
eval('2 + 2') # high risk
|
|
"""
|
|
results1 = detectGPT(content)
|
|
self.assertEqual(len(results1["high"]), 3)
|
|
|
|
def test_gpt_no_risk_detection(self):
|
|
"""
|
|
if os.getenv("BAIDU_API_KEY") is None or os.getenv("BAIDU_SECRET_KEY") is None:
|
|
warnings.warn("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set, test skipped.", UserWarning)
|
|
self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
|
|
"""
|
|
content = """a = 10
|
|
b = a + 5
|
|
print('This should not be detected as risky.')
|
|
"""
|
|
results2 = detectGPT(content)
|
|
self.assertEqual(len(results2["high"]), 0)
|
|
self.assertEqual(len(results2["medium"]), 0)
|
|
self.assertEqual(len(results2["low"]), 0)
|
|
|
|
def test_gpt_env_no_set(self):
|
|
"""
|
|
if os.getenv("BAIDU_API_KEY") is not None or os.getenv("BAIDU_SECRET_KEY") is not None:
|
|
self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is set")
|
|
"""
|
|
content = "print('test test')"
|
|
with self.assertRaises(ValueError):
|
|
detectGPT(content)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|