Spaces:
Running
Running
File size: 5,041 Bytes
168ae1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | # Unit tests for the AI Code Security Scanner
import unittest
import sys
import os
# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from rule_detector import RuleBasedCodeDetector
from combined_detector import CombinedCodeDetector
from fix_generator import FixSuggestionGenerator
class TestRuleBasedDetector(unittest.TestCase):
def setUp(self):
self.detector = RuleBasedCodeDetector()
def test_sql_injection(self):
code = """query = f"SELECT * FROM users WHERE id = {user_id}" """
result = self.detector.analyze(code)
self.assertGreater(result["issue_count"], 0)
self.assertTrue(any(i["type"] == "sql_injection" for i in result["issues"]))
def test_hardcoded_secret(self):
code = """api_key = "sk_live_1234567890" """
result = self.detector.analyze(code)
self.assertTrue(any(i["type"] == "hardcoded_secret" for i in result["issues"]))
def test_safe_code(self):
code = """def safe_function():\n return 42"""
result = self.detector.analyze(code)
self.assertEqual(result["issue_count"], 0)
self.assertGreater(result["security_score"], 80)
def test_syntax_error(self):
code = """def test()\n print("hello")""" # Missing colon
result = self.detector.analyze(code)
self.assertGreater(result["issue_count"], 0)
self.assertTrue(any(i["type"] == "syntax_error" for i in result["issues"]))
class TestCombinedDetector(unittest.TestCase):
def setUp(self):
self.detector = CombinedCodeDetector()
def test_combined_analysis(self):
code = """query = f"SELECT * FROM users WHERE id = {user_id}" """
result = self.detector.combined_analysis(code)
self.assertIn("security_score", result)
self.assertIn("issues", result)
self.assertIn("ml_analysis", result)
def test_vulnerable_vs_safe(self):
vulnerable = """api_key = "sk_test_1234567890" """
safe = """api_key = os.getenv("API_KEY") """
vuln_result = self.detector.combined_analysis(vulnerable)
safe_result = self.detector.combined_analysis(safe)
# Vulnerable code should have lower score
self.assertLess(vuln_result["security_score"], safe_result["security_score"])
class TestFixGenerator(unittest.TestCase):
def setUp(self):
self.generator = FixSuggestionGenerator()
def test_get_fixes(self):
fixes = self.generator.get_fixes(
"""query = f"SELECT * FROM users WHERE id = {user_id}" """,
"sql_injection"
)
self.assertGreater(len(fixes), 0)
self.assertIsInstance(fixes, list)
def test_generate_fix_patch(self):
code = """def test():\n query = f"SELECT * FROM users WHERE id = {user_id}"\n return query"""
patch = self.generator.generate_fix_patch(code, 2, "sql_injection")
self.assertIn("fix_suggestions", patch)
self.assertIn("vulnerable_line", patch)
class TestIntegration(unittest.TestCase):
# Integration tests
def test_end_to_end(self):
# Test that all components work together
detector = CombinedCodeDetector()
generator = FixSuggestionGenerator()
test_code = """def insecure():\n query = f"SELECT * FROM users WHERE id = {user_id}"\n key = "secret123"\n return query"""
# Detect issues
analysis = detector.combined_analysis(test_code)
# Generate fixes for first issue
if analysis["issues"]:
issue = analysis["issues"][0]
fixes = generator.get_fixes(issue["message"], issue["type"])
self.assertGreater(len(fixes), 0)
self.assertGreater(analysis["issue_count"], 0)
def run_tests():
# Run all tests and print results
print("🧪 Running AI Code Security Scanner Tests")
print("="*50)
# Create test suite
loader = unittest.TestLoader()
suite = unittest.TestSuite()
suite.addTests(loader.loadTestsFromTestCase(TestRuleBasedDetector))
suite.addTests(loader.loadTestsFromTestCase(TestCombinedDetector))
suite.addTests(loader.loadTestsFromTestCase(TestFixGenerator))
suite.addTests(loader.loadTestsFromTestCase(TestIntegration))
# Run tests
runner = unittest.TextTestRunner(verbosity=2)
result = runner.run(suite)
print("\n" + "="*50)
print(f"📊 Test Results: {result.testsRun} tests run")
print(f"✅ Passed: {result.testsRun - len(result.failures) - len(result.errors)}")
print(f"❌ Failed: {len(result.failures)}")
print(f"⚠️ Errors: {len(result.errors)}")
if result.wasSuccessful():
print("\n🎉 All tests passed!")
else:
print("\n🔧 Some tests failed. Review and fix.")
return result.wasSuccessful()
if __name__ == "__main__":
success = run_tests()
sys.exit(0 if success else 1) |