🐍 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
165
scripts/fix_async_audit_queue_tests_v2.py
Executable file
165
scripts/fix_async_audit_queue_tests_v2.py
Executable file
@@ -0,0 +1,165 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Fixed migration script for AsyncAuditQueue API changes.
|
||||
"""
|
||||
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
def fix_file_content(content: str) -> str:
|
||||
"""Apply all fixes to file content."""
|
||||
|
||||
# Fix 1: AsyncAuditQueue::new() - add parameters and .await
|
||||
# Pattern: AsyncAuditQueue::new(wal_path_expr)
|
||||
# Replace: AsyncAuditQueue::new(wal_path_expr, Arc::clone(&pool), 100, 100).await.expect("...")
|
||||
content = re.sub(
|
||||
r'AsyncAuditQueue::new\(([^)]+)\)(?!\s*\.await)',
|
||||
r'AsyncAuditQueue::new(\1, Arc::clone(&pool), 100, 100).await.expect("Failed to create AsyncAuditQueue")',
|
||||
content
|
||||
)
|
||||
|
||||
# Fix 2: Remove Arc::new() wrapper from AsyncAuditQueue (it returns Arc now)
|
||||
content = re.sub(
|
||||
r'Arc::new\(AsyncAuditQueue::new\(([^)]+), Arc::clone\(&pool\), 100, 100\)\.await\.expect\("Failed to create AsyncAuditQueue"\)\)',
|
||||
r'AsyncAuditQueue::new(\1, Arc::clone(&pool), 100, 100).await.expect("Failed to create AsyncAuditQueue")',
|
||||
content
|
||||
)
|
||||
|
||||
# Fix 3: AuditTrailEngine::new() - add pool and wal_path parameters
|
||||
# Only fix if it doesn't already have .await
|
||||
content = re.sub(
|
||||
r'AuditTrailEngine::new\(([^,)]+)\)(?!\s*\.await)',
|
||||
r'AuditTrailEngine::new(\1, Arc::clone(&pool), wal_path.clone()).await.expect("Failed to create AuditTrailEngine")',
|
||||
content
|
||||
)
|
||||
|
||||
# Fix 4: PostgresPool::new() with 2 args -> PostgresPool::connect() with 1 arg
|
||||
content = re.sub(
|
||||
r'PostgresPool::new\(([^,)]+),\s*\d+\)',
|
||||
r'PostgresPool::connect(\1).await.expect("Failed to connect to PostgreSQL")',
|
||||
content
|
||||
)
|
||||
|
||||
# Fix 5: Add missing pool setup at start of test functions if needed
|
||||
# This is done more carefully per-function
|
||||
lines = content.split('\n')
|
||||
result = []
|
||||
i = 0
|
||||
|
||||
while i < len(lines):
|
||||
line = lines[i]
|
||||
result.append(line)
|
||||
|
||||
# Check if this is an async test function
|
||||
if 'async fn test_' in line or 'async fn bench_' in line:
|
||||
# Look for opening brace
|
||||
j = i + 1
|
||||
while j < len(lines) and '{' not in lines[j]:
|
||||
result.append(lines[j])
|
||||
j += 1
|
||||
|
||||
if j < len(lines):
|
||||
result.append(lines[j]) # Add the line with {
|
||||
|
||||
# Check if next few lines already have pool setup
|
||||
has_pool = False
|
||||
for k in range(j + 1, min(j + 5, len(lines))):
|
||||
if 'let pool =' in lines[k] or 'setup_test_postgres_pool' in lines[k]:
|
||||
has_pool = True
|
||||
break
|
||||
|
||||
# Check if function uses AsyncAuditQueue or AuditTrailEngine
|
||||
needs_pool = False
|
||||
for k in range(j + 1, min(j + 30, len(lines))):
|
||||
if 'AsyncAuditQueue::new' in lines[k] or 'AuditTrailEngine::new' in lines[k]:
|
||||
needs_pool = True
|
||||
break
|
||||
|
||||
# Add pool setup if needed
|
||||
if needs_pool and not has_pool:
|
||||
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('')
|
||||
|
||||
i = j
|
||||
|
||||
i += 1
|
||||
|
||||
content = '\n'.join(result)
|
||||
|
||||
# Fix 6: Remove old AuditTrailConfig fields
|
||||
old_fields_patterns = [
|
||||
r'\s+enabled:\s*\w+,?\s*\n',
|
||||
r'\s+compression_algorithm:\s*CompressionAlgorithm::\w+,?\s*\n',
|
||||
r'\s+encryption_algorithm:\s*EncryptionAlgorithm::\w+,?\s*\n',
|
||||
r'\s+encryption_key:\s*vec!\[.*?\],?\s*\n',
|
||||
r'\s+postgres_pool:\s*[^,\n]+,?\s*\n',
|
||||
r'\s+file_path:\s*[^,\n]+,?\s*\n',
|
||||
r'\s+enable_checksums:\s*\w+,?\s*\n',
|
||||
r'\s+enable_tamper_detection:\s*\w+,?\s*\n',
|
||||
r'\s+enable_best_execution_tracking:\s*\w+,?\s*\n',
|
||||
r'\s+enable_mifid_reporting:\s*\w+,?\s*\n',
|
||||
]
|
||||
|
||||
for pattern in old_fields_patterns:
|
||||
content = re.sub(pattern, '', content)
|
||||
|
||||
# Fix 7: Fix enum variants
|
||||
content = content.replace('EncryptionAlgorithm::Aes256Gcm', 'EncryptionAlgorithm::AES256GCM')
|
||||
|
||||
# Fix 8: Remove duplicate PostgresPool::connect().await.expect().await patterns
|
||||
content = re.sub(
|
||||
r'PostgresPool::connect\(([^)]+)\)\.await\.expect\("([^"]+)"\)\.await',
|
||||
r'PostgresPool::connect(\1).await.expect("\2")',
|
||||
content
|
||||
)
|
||||
|
||||
return content
|
||||
|
||||
def process_file(file_path: Path) -> bool:
|
||||
"""Process a single test file."""
|
||||
print(f"Processing {file_path.name}...")
|
||||
|
||||
try:
|
||||
original = file_path.read_text()
|
||||
fixed = fix_file_content(original)
|
||||
|
||||
if fixed != original:
|
||||
file_path.write_text(fixed)
|
||||
print(f" ✅ Fixed {file_path.name}")
|
||||
return True
|
||||
else:
|
||||
print(f" ⏭️ No changes for {file_path.name}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f" ❌ Error: {e}")
|
||||
return False
|
||||
|
||||
def main():
|
||||
test_dir = Path(__file__).parent.parent / 'trading_engine' / 'tests'
|
||||
|
||||
if not test_dir.exists():
|
||||
print(f"Error: {test_dir} not found")
|
||||
return 1
|
||||
|
||||
test_files = sorted(test_dir.glob('*.rs'))
|
||||
|
||||
if not test_files:
|
||||
print(f"Error: No test files in {test_dir}")
|
||||
return 1
|
||||
|
||||
print(f"Found {len(test_files)} test files\n")
|
||||
|
||||
updated = sum(1 for f in test_files if process_file(f))
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Updated {updated}/{len(test_files)} files")
|
||||
print(f"{'='*60}")
|
||||
|
||||
return 0
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user