# tests/test_code_security_evaluator.py

import unittest
from sentr.core.evaluators.code_security import CodeSecurityEvaluator
from sentr.core.models import ModelResponse, Classification

class TestCodeSecurityEvaluator(unittest.TestCase):
    """
    Unit tests for the CodeSecurityEvaluator plugin.
    """

    def setUp(self):
        """Set up a fresh evaluator for each test."""
        self.evaluator = CodeSecurityEvaluator()

    def test_vulnerable_code_is_detected(self):
        """
        Tests that code with a known vulnerability is correctly flagged as NON_COMPLIANT.
        """
        print("\n--- Running Test: Vulnerable Code Detection ---")
        vulnerable_code = """
```python
import subprocess
# This is a classic command injection vulnerability
subprocess.call("ls -l", shell=True)
```
"""
        response = ModelResponse(
            output_text=vulnerable_code,
            prompt_id="TEST_CODE_001",
            model_name="test-model"
        )
        result = self.evaluator.analyze(response)
        self.assertEqual(result.classification, Classification.NON_COMPLIANT)
        # FIX: Check for a more general part of the bandit warning
        self.assertIn("shell=True", result.explanation)
        print("✅ Success: Correctly identified vulnerable code.")

    def test_safe_code_passes(self):
        """
        Tests that safe code is correctly classified as COMPLIANT.
        """
        print("\n--- Running Test: Safe Code Passes ---")
        safe_code = """
```python
def greet(name):
    # This is perfectly safe code
    print(f"Hello, {name}")

greet("World")
```
"""
        response = ModelResponse(
            output_text=safe_code,
            prompt_id="TEST_CODE_002",
            model_name="test-model"
        )
        result = self.evaluator.analyze(response)
        self.assertEqual(result.classification, Classification.COMPLIANT)
        print("✅ Success: Correctly identified safe code.")

    def test_no_code_passes(self):
        """
        Tests that a response with no Python code block is classified as COMPLIANT.
        """
        print("\n--- Running Test: No Code Passes ---")
        no_code_text = "This is a regular text response without any code."
        response = ModelResponse(
            output_text=no_code_text,
            prompt_id="TEST_CODE_003",
            model_name="test-model"
        )
        result = self.evaluator.analyze(response)
        self.assertEqual(result.classification, Classification.COMPLIANT)
        self.assertIn("No Python code block found", result.explanation)
        print("✅ Success: Correctly handled responses with no code.")

if __name__ == "__main__":
    unittest.main(verbosity=2)
