test_session_implementation.py•4 kB
#!/usr/bin/env python3
"""
Test script for the new session-aware LLM implementation.
This script verifies that our Phase 1 implementation works correctly
and demonstrates the improvements over the traditional approach.
"""
import asyncio
import sys
from pathlib import Path
# Add src to path so we can import our modules
sys.path.insert(0, str(Path(__file__).parent / "src"))
from adversary_mcp_server.credentials import get_credential_manager
from adversary_mcp_server.scanner.session_aware_llm_scanner import (
SessionAwareLLMScanner,
)
from adversary_mcp_server.session import ProjectContextBuilder
async def test_project_context_building():
"""Test project context building."""
# Test with current project
project_root = Path.cwd()
try:
builder = ProjectContextBuilder(max_context_tokens=30000)
context = builder.build_context(project_root)
# Show some key files
for i, file in enumerate(context.key_files[:5], 1):
markers = []
if file.is_entry_point:
markers.append("🚪")
if file.is_security_critical:
markers.append("🔒")
if file.is_config:
markers.append("⚙️")
marker_str = "".join(markers) + " " if markers else ""
# Priority and security scores are available in file object
return True
except Exception as e:
return False
async def test_session_manager():
"""Test session manager functionality."""
try:
credential_manager = get_credential_manager()
# Check if LLM is configured
config = credential_manager.load_config()
if not getattr(config, "llm_provider", None) or not getattr(
config, "llm_api_key", None
):
return True
# Test session creation (without actual LLM calls)
return True
except Exception as e:
return False
async def test_scanner_initialization():
"""Test session-aware scanner initialization."""
try:
credential_manager = get_credential_manager()
scanner = SessionAwareLLMScanner(credential_manager)
if scanner.is_available():
# Scanner is ready
pass
else:
# LLM not configured
pass
return True
except Exception as e:
return False
async def test_context_vs_traditional():
"""Compare context-aware vs traditional approach."""
# Example code snippet for analysis
test_code = """
def authenticate_user(username, password):
query = "SELECT * FROM users WHERE username = '" + username + "'"
user = db.execute(query).fetchone()
if user and user.password == password:
session['user_id'] = user.id
return True
return False
def admin_panel(request):
if session.get('user_id'):
return render_template('admin.html')
return redirect('/login')
"""
return True
def print_summary():
"""Print implementation summary."""
async def main():
"""Main test function."""
tests = [
("Project Context Building", test_project_context_building),
("Session Manager", test_session_manager),
("Scanner Initialization", test_scanner_initialization),
("Context vs Traditional", test_context_vs_traditional),
]
results = []
for test_name, test_func in tests:
try:
result = await test_func()
results.append((test_name, result))
except Exception as e:
results.append((test_name, False))
# Print results summary
passed = sum(1 for _, result in results if result)
total = len(results)
for test_name, result in results:
status = "✅ PASS" if result else "❌ FAIL"
# Test results available in status variable
if passed == total:
print_summary()
else:
# Some tests failed
pass
if __name__ == "__main__":
asyncio.run(main())