backend/ai/tests/validation_test.py

"""
Simple validation test for AegisLM Red Team Engine.
This test validates the basic structure without external dependencies.
"""

import sys
import os

# Add the AI package root to Python path
ai_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ai_root)

# Now we can import with relative paths from the AI root


def test_basic_imports():
    """Test that basic modules can be imported."""
    try:
        # Test engine imports
        from engines.attack_engine import AttackEngine
        from engines.mutation_engine import MutationEngine
        print("✓ Engine imports successful")
        
        # Test agent imports
        from agents.attacker_agent import AttackerAgent
        from agents.evaluator_agent import EvaluatorAgent
        from agents.coordinator_agent import CoordinatorAgent
        print("✓ Agent imports successful")
        
        # Test pipeline imports
        from pipelines.redteam_pipeline import RedTeamPipeline
        print("✓ Pipeline imports successful")
        
        return True
    except ImportError as e:
        print(f"✗ Import error: {e}")
        return False


def test_attack_engine_basic():
    """Test basic attack engine functionality."""
    try:
        from engines.attack_engine import AttackEngine
        from models.attack_models import AttackType
        
        engine = AttackEngine()
        attack_types = engine.get_supported_attack_types()
        
        assert len(attack_types) == 4
        assert AttackType.JAILBREAK in attack_types
        assert AttackType.PROMPT_INJECTION in attack_types
        assert AttackType.TOXICITY_TRIGGER in attack_types
        assert AttackType.HALLUCINATION_TRAP in attack_types
        
        print("✓ Attack engine basic functionality works")
        return True
    except Exception as e:
        print(f"✗ Attack engine test failed: {e}")
        return False


def test_mutation_engine_basic():
    """Test basic mutation engine functionality."""
    try:
        from engines.mutation_engine import MutationEngine
        from models.attack_models import Attack, AttackType, SeverityLevel
        
        engine = MutationEngine()
        strategies = engine.get_available_strategies()
        
        assert len(strategies) == 4
        assert "obfuscation" in strategies
        assert "roleplay" in strategies
        assert "indirect_phrasing" in strategies
        assert "escalation" in strategies
        
        print("✓ Mutation engine basic functionality works")
        return True
    except Exception as e:
        print(f"✗ Mutation engine test failed: {e}")
        return False


def test_pipeline_structure():
    """Test that pipeline structure is correct."""
    try:
        from pipelines.redteam_pipeline import RedTeamPipeline
        
        pipeline = RedTeamPipeline()
        assert hasattr(pipeline, 'run_redteam_pipeline')
        assert hasattr(pipeline, 'run_targeted_assessment')
        assert hasattr(pipeline, 'run_iterative_improvement')
        
        print("✓ Pipeline structure is correct")
        return True
    except Exception as e:
        print(f"✗ Pipeline structure test failed: {e}")
        return False


def test_file_structure():
    """Test that all required files exist."""
    required_files = [
        'models/attack_models.py',
        'models/result_models.py', 
        'models/scoring_models.py',
        'engines/attack_engine.py',
        'engines/mutation_engine.py',
        'engines/scoring_engine.py',
        'agents/attacker_agent.py',
        'agents/evaluator_agent.py',
        'agents/coordinator_agent.py',
        'pipelines/redteam_pipeline.py',
        'utils/prompt_templates.py'
    ]
    
    base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    
    for file_path in required_files:
        full_path = os.path.join(base_path, file_path)
        if not os.path.exists(full_path):
            print(f"✗ Missing file: {file_path}")
            return False
    
    print("✓ All required files exist")
    return True


def test_model_interface():
    """Test that model interface is properly defined."""
    try:
        from agents.attacker_agent import ModelInterface
        
        # Check that ModelInterface is a proper Protocol
        assert hasattr(ModelInterface, 'generate_response')
        
        # Test with a mock implementation
        class MockModel:
            def generate_response(self, prompt: str) -> str:
                return f"Response to: {prompt}"
        
        mock = MockModel()
        response = mock.generate_response("test prompt")
        assert response == "Response to: test prompt"
        
        print("✓ Model interface is properly defined")
        return True
    except Exception as e:
        print(f"✗ Model interface test failed: {e}")
        return False


def main():
    """Run all validation tests."""
    print("🚀 Running AegisLM Red Team Engine Validation Tests")
    print("=" * 60)
    
    tests = [
        ("File Structure", test_file_structure),
        ("Basic Imports", test_basic_imports),
        ("Attack Engine", test_attack_engine_basic),
        ("Mutation Engine", test_mutation_engine_basic),
        ("Pipeline Structure", test_pipeline_structure),
        ("Model Interface", test_model_interface)
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        print(f"\n📋 Testing {test_name}...")
        if test_func():
            passed += 1
        else:
            print(f"❌ {test_name} failed")
    
    print("\n" + "=" * 60)
    print(f"📊 Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All validation tests passed!")
        print("\n📝 Next Steps:")
        print("1. Install dependencies: pip install -r requirements.txt")
        print("2. Run full test suite: python tests/test_redteam_engine.py")
        print("3. Check README.md for usage examples")
        return True
    else:
        print("⚠️  Some tests failed. Please check the errors above.")
        return False


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)