backend/ai/setup_tests.py

"""
Setup script for AegisLM Red Team Engine testing.

Helps users configure environment and install dependencies.
"""

import os
import sys
import subprocess
from pathlib import Path

def install_dependencies():
    """Install test dependencies."""
    print("📦 Installing test dependencies...")
    try:
        subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements-test.txt"], 
                      check=True, capture_output=True, text=True)
        print("✅ Dependencies installed successfully")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ Failed to install dependencies: {e}")
        print(f"Output: {e.stdout}")
        print(f"Error: {e.stderr}")
        return False

def setup_env_file():
    """Set up .env file for API keys."""
    env_file = Path(".env")
    env_example = Path(".env.example")
    
    if env_file.exists():
        print("✅ .env file already exists")
        print("📝 Edit .env file to add your API keys:")
        print("   GROK_API_KEY=your_grok_api_key")
        print("   MISTRAL_API_KEY=your_mistral_api_key")
        print("   RUN_INTEGRATION_TESTS=true")
        return True
    
    if env_example.exists():
        # Copy example to .env
        with open(env_example, 'r') as f:
            content = f.read()
        
        with open(env_file, 'w') as f:
            f.write(content)
        
        print("✅ Created .env file from .env.example")
        print("📝 Edit .env file to add your API keys:")
        print("   GROK_API_KEY=your_grok_api_key")
        print("   MISTRAL_API_KEY=your_mistral_api_key")
        print("   RUN_INTEGRATION_TESTS=true")
        return True
    else:
        # Create basic .env file
        env_content = """# Environment Variables for AegisLM Red Team Engine
# Fill in your actual API keys below

# Integration Test Configuration
RUN_INTEGRATION_TESTS=false

# Real LLM API Keys (for integration tests only)
GROK_API_KEY=
MISTRAL_API_KEY=

# Optional: Additional Configuration
API_TIMEOUT_SECONDS=30
MAX_RETRIES=3
RATE_LIMIT_DELAY=1.0
"""
        
        with open(env_file, 'w') as f:
            f.write(env_content)
        
        print("✅ Created .env file")
        print("📝 Edit .env file to add your API keys:")
        print("   GROK_API_KEY=your_grok_api_key")
        print("   MISTRAL_API_KEY=your_mistral_api_key")
        print("   RUN_INTEGRATION_TESTS=true")
        return True

def check_setup():
    """Check if setup is complete."""
    print("🔍 Checking setup...")
    
    # Check dependencies
    try:
        import dotenv
        print("✅ python-dotenv installed")
    except ImportError:
        print("❌ python-dotenv not installed")
        return False
    
    # Check .env file
    if Path(".env").exists():
        print("✅ .env file exists")
    else:
        print("❌ .env file not found")
        return False
    
    # Check API keys
    from tests import get_test_config
    config = get_test_config()
    
    print(f"🔑 API Status:")
    print(f"  GROK_API_KEY: {'✅ Set' if config['has_grok_key'] else '❌ Not set'}")
    print(f"  MISTRAL_API_KEY: {'✅ Set' if config['has_mistral_key'] else '❌ Not set'}")
    print(f"  Integration Tests: {'✅ Enabled' if config['run_integration'] else '❌ Disabled'}")
    
    return True

def main():
    """Main setup function."""
    print("🚀 AegisLM Red Team Engine - Test Setup")
    print("=" * 50)
    
    # Install dependencies
    if not install_dependencies():
        print("\n❌ Setup failed during dependency installation")
        return False
    
    # Set up .env file
    if not setup_env_file():
        print("\n❌ Setup failed during .env file creation")
        return False
    
    # Check setup
    if not check_setup():
        print("\n❌ Setup verification failed")
        return False
    
    print("\n✅ Setup completed successfully!")
    print("\n📝 Next steps:")
    print("1. Edit .env file to add your API keys")
    print("2. Set RUN_INTEGRATION_TESTS=true to enable integration tests")
    print("3. Run tests with: python run_tests.py")
    print("4. For integration tests: python run_tests.py --integration")
    
    return True

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