🐍 Wave 112: Python fix scripts for audit tests
- 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
This commit is contained in:
jgrusewski
237
scripts/fix_async_audit_queue_tests.py
Executable file
237
scripts/fix_async_audit_queue_tests.py
Executable file
@@ -0,0 +1,237 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user