#!/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())