a7973d4bf7
- fix_all_audit_tests.py: Comprehensive audit test fixes - fix_async_audit_queue_tests.py: Async queue test fixes (v1) - fix_async_audit_queue_tests_v2.py: Async queue test fixes (v2) - fix_audit_compliance.py: Compliance test fixes
238 lines
8.5 KiB
Python
Executable File
238 lines
8.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Migration script to fix AsyncAuditQueue test compilation errors.
|
|
|
|
This script updates all test files to use the new 4-parameter async AsyncAuditQueue API:
|
|
- Old: Arc::new(AsyncAuditQueue::new(wal_path))
|
|
- New: AsyncAuditQueue::new(wal_path, pool, batch_size, flush_interval).await?
|
|
|
|
Changes applied:
|
|
1. Add .await? to AsyncAuditQueue::new() calls
|
|
2. Add 3 new parameters (pool, batch_size, flush_interval)
|
|
3. Remove Arc::new() wrapping (returns Arc internally)
|
|
4. Fix function signatures to be async
|
|
5. Fix AuditTrailEngine::new() calls (add pool and wal_path parameters)
|
|
"""
|
|
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
def fix_async_audit_queue_new(content: str) -> str:
|
|
"""Fix AsyncAuditQueue::new() calls to use new 4-parameter async API."""
|
|
|
|
# Pattern 1: Arc::new(AsyncAuditQueue::new(wal_path))
|
|
# Replace with: AsyncAuditQueue::new(wal_path, Arc::clone(&pool), 100, 100).await?
|
|
pattern1 = r'Arc::new\(AsyncAuditQueue::new\(([^)]+)\)\)'
|
|
replacement1 = r'AsyncAuditQueue::new(\1, Arc::clone(&pool), 100, 100).await.expect("Failed to create AsyncAuditQueue")'
|
|
content = re.sub(pattern1, replacement1, content)
|
|
|
|
# Pattern 2: AsyncAuditQueue::new(wal_path) without Arc::new
|
|
# Replace with: AsyncAuditQueue::new(wal_path, Arc::clone(&pool), 100, 100).await?
|
|
pattern2 = r'AsyncAuditQueue::new\(([^,)]+)\)(?!\s*\.await)'
|
|
replacement2 = r'AsyncAuditQueue::new(\1, Arc::clone(&pool), 100, 100).await.expect("Failed to create AsyncAuditQueue")'
|
|
content = re.sub(pattern2, replacement2, content)
|
|
|
|
return content
|
|
|
|
def fix_audit_trail_engine_new(content: str) -> str:
|
|
"""Fix AuditTrailEngine::new() calls to use new 3-parameter async API."""
|
|
|
|
# Pattern: AuditTrailEngine::new(config) -> AuditTrailEngine::new(config, pool, wal_path).await?
|
|
pattern = r'AuditTrailEngine::new\(([^)]+)\)(?!\s*\.await)'
|
|
|
|
def replacement(match):
|
|
config = match.group(1).strip()
|
|
# If only config, add pool and wal_path
|
|
if ',' not in config:
|
|
return f'AuditTrailEngine::new({config}, Arc::clone(&pool), wal_path.clone()).await.expect("Failed to create AuditTrailEngine")'
|
|
return match.group(0)
|
|
|
|
content = re.sub(pattern, replacement, content)
|
|
return content
|
|
|
|
def add_async_to_test_functions(content: str) -> str:
|
|
"""Add async to test functions that use .await."""
|
|
|
|
# Pattern: #[test] or #[tokio::test] followed by fn function_name()
|
|
# If function body contains .await, ensure #[tokio::test]
|
|
lines = content.split('\n')
|
|
result = []
|
|
i = 0
|
|
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
|
|
# Check if this is a test function
|
|
if '#[test]' in line or '#[tokio::test]' in line:
|
|
# Look ahead to find the function body
|
|
fn_start = i + 1
|
|
while fn_start < len(lines) and 'fn ' not in lines[fn_start]:
|
|
fn_start += 1
|
|
|
|
if fn_start < len(lines):
|
|
# Check if function uses .await
|
|
fn_line = lines[fn_start]
|
|
|
|
# Find function body end
|
|
brace_count = 0
|
|
fn_body_lines = []
|
|
j = fn_start
|
|
started = False
|
|
|
|
while j < len(lines):
|
|
if '{' in lines[j]:
|
|
started = True
|
|
brace_count += lines[j].count('{')
|
|
if '}' in lines[j]:
|
|
brace_count -= lines[j].count('}')
|
|
|
|
fn_body_lines.append(lines[j])
|
|
j += 1
|
|
|
|
if started and brace_count == 0:
|
|
break
|
|
|
|
fn_body = '\n'.join(fn_body_lines)
|
|
|
|
# If function uses .await, ensure it's async and uses #[tokio::test]
|
|
if '.await' in fn_body:
|
|
# Replace #[test] with #[tokio::test]
|
|
if '#[test]' in line:
|
|
line = line.replace('#[test]', '#[tokio::test]')
|
|
|
|
# Make function async if not already
|
|
if 'async fn ' not in fn_line:
|
|
lines[fn_start] = fn_line.replace('fn ', 'async fn ')
|
|
|
|
result.append(line)
|
|
i += 1
|
|
|
|
return '\n'.join(result)
|
|
|
|
def add_pool_setup(content: str) -> str:
|
|
"""Add PostgreSQL pool setup to test functions that need it."""
|
|
|
|
lines = content.split('\n')
|
|
result = []
|
|
|
|
for i, line in enumerate(lines):
|
|
result.append(line)
|
|
|
|
# If this is a test function that uses AsyncAuditQueue or AuditTrailEngine
|
|
if ('async fn test_' in line or 'async fn bench_' in line) and i + 1 < len(lines):
|
|
# Look ahead for AsyncAuditQueue or AuditTrailEngine usage
|
|
body_start = i + 1
|
|
while body_start < len(lines) and '{' not in lines[body_start]:
|
|
body_start += 1
|
|
|
|
if body_start < len(lines):
|
|
# Check next 50 lines for usage
|
|
check_lines = '\n'.join(lines[body_start:min(body_start + 50, len(lines))])
|
|
|
|
if ('AsyncAuditQueue::new' in check_lines or 'AuditTrailEngine::new' in check_lines) and 'setup_test_postgres_pool' not in check_lines:
|
|
# Add pool setup after opening brace
|
|
result.append(' // Setup PostgreSQL pool for audit persistence')
|
|
result.append(' let pool = setup_test_postgres_pool().await;')
|
|
result.append(' let wal_path = std::env::temp_dir().join(format!("foxhunt_audit_test_{}.wal", uuid::Uuid::new_v4()));')
|
|
result.append('')
|
|
|
|
return '\n'.join(result)
|
|
|
|
def fix_audit_config_fields(content: str) -> str:
|
|
"""Fix AuditTrailConfig field names that have changed."""
|
|
|
|
# Remove old fields that no longer exist
|
|
old_fields = [
|
|
r'\s*enabled:\s*true,?\n',
|
|
r'\s*compression_algorithm:\s*CompressionAlgorithm::\w+,?\n',
|
|
r'\s*encryption_algorithm:\s*EncryptionAlgorithm::\w+,?\n',
|
|
r'\s*encryption_key:\s*vec!\[.*?\],?\n',
|
|
r'\s*postgres_pool:\s*\w+,?\n',
|
|
r'\s*file_path:\s*\w+,?\n',
|
|
r'\s*enable_checksums:\s*\w+,?\n',
|
|
r'\s*enable_tamper_detection:\s*\w+,?\n',
|
|
r'\s*enable_best_execution_tracking:\s*\w+,?\n',
|
|
]
|
|
|
|
for pattern in old_fields:
|
|
content = re.sub(pattern, '', content)
|
|
|
|
# Fix enum variant names
|
|
content = content.replace('EncryptionAlgorithm::Aes256Gcm', 'EncryptionAlgorithm::AES256GCM')
|
|
|
|
return content
|
|
|
|
def fix_postgres_pool_new(content: str) -> str:
|
|
"""Fix PostgresPool::new() calls - it takes only 1 argument (connection string)."""
|
|
|
|
# Pattern: PostgresPool::new(&database_url, pool_size)
|
|
# Replace with: PostgresPool::connect(&database_url).await?
|
|
pattern = r'PostgresPool::new\(([^,)]+)(?:,\s*\d+)?\)'
|
|
replacement = r'PostgresPool::connect(\1).await.expect("Failed to connect to PostgreSQL")'
|
|
content = re.sub(pattern, replacement, content)
|
|
|
|
return content
|
|
|
|
def process_file(file_path: Path) -> bool:
|
|
"""Process a single test file."""
|
|
|
|
print(f"Processing {file_path}...")
|
|
|
|
try:
|
|
content = file_path.read_text()
|
|
original = content
|
|
|
|
# Apply all fixes
|
|
content = fix_async_audit_queue_new(content)
|
|
content = fix_audit_trail_engine_new(content)
|
|
content = add_async_to_test_functions(content)
|
|
content = add_pool_setup(content)
|
|
content = fix_audit_config_fields(content)
|
|
content = fix_postgres_pool_new(content)
|
|
|
|
# Only write if changed
|
|
if content != original:
|
|
file_path.write_text(content)
|
|
print(f" ✅ Updated {file_path.name}")
|
|
return True
|
|
else:
|
|
print(f" ⏭️ No changes needed for {file_path.name}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f" ❌ Error processing {file_path}: {e}")
|
|
return False
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
|
|
# Find all test files in trading_engine/tests/
|
|
test_dir = Path(__file__).parent.parent / 'trading_engine' / 'tests'
|
|
|
|
if not test_dir.exists():
|
|
print(f"Error: Test directory not found: {test_dir}")
|
|
sys.exit(1)
|
|
|
|
test_files = list(test_dir.glob('*.rs'))
|
|
|
|
if not test_files:
|
|
print(f"Error: No test files found in {test_dir}")
|
|
sys.exit(1)
|
|
|
|
print(f"Found {len(test_files)} test files to process\n")
|
|
|
|
updated_count = 0
|
|
for test_file in sorted(test_files):
|
|
if process_file(test_file):
|
|
updated_count += 1
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Summary: Updated {updated_count} of {len(test_files)} files")
|
|
print(f"{'='*60}")
|
|
|
|
return 0
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main())
|