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()