Spaces:
Running
Running
| # 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) |