sangge/BackDoorBuster
1
0
Fork 0
Code Issues Pull Requests Projects 1 Activity

update：修改国内gpt调用

Browse Source
This commit is contained in:
ccyj
2024-05-24 20:27:18 +08:00
parent f0e2251dc0
commit b1bc566c09
2 changed files with 58 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
tests/test_CN_GPT_detection.py
View File

@@ -1,49 +1,47 @@
import unittest
import warnings
import os
import json
from detection.cngptdetection import detectGPT # 导入调用百度 ai 模型的函数
from detection.cngptdetection import detectGPT
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:
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)
classified_results = json.loads(results1)
self.assertEqual(len(classified_results["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)
classified_results = json.loads(results2)
self.assertEqual(len(classified_results["high"]), 0)
self.assertEqual(len(classified_results["medium"]), 0)
self.assertEqual(len(classified_results["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()