use probe_code::search::search_runner::{perform_probe, search_with_structured_patterns};
use probe_code::search::SearchOptions;
use std::fs;
use std::path::{Path, PathBuf};
use std::time::Instant;
use tempfile::tempdir;
fn create_test_files(dir: &Path, count: usize, lines_per_file: usize) -> Vec<PathBuf> {
let mut paths = Vec::new();
for i in 0..count {
let file_path = dir.join(format!("test_file_{i}.rs"));
let mut content = format!("// Test file {i}\n\n");
for j in 0..lines_per_file / 10 {
content.push_str(&format!(
"/// Function documentation for function_{}_{}
fn function_{}_{}() {{
// Initialize variables
let search_term_alpha = {};
let search_term_beta = {};
// Process data
println!(\"Processing data with search_term_gamma\");
// Return result
search_term_alpha + search_term_beta
}}
",
i,
j,
i,
j,
j * 2,
j * 3
));
}
fs::write(&file_path, content).unwrap();
paths.push(file_path);
}
paths
}
#[test]
fn test_parallel_file_search() {
let temp_dir = tempdir().unwrap();
let base_path = temp_dir.path();
let file_count = 50;
let lines_per_file = 100;
let _file_paths = create_test_files(base_path, file_count, lines_per_file);
let queries = vec![
"search_term_alpha".to_string(),
"search_term_beta".to_string(),
];
let custom_ignores = Vec::new();
let options = SearchOptions {
path: base_path,
queries: &queries,
files_only: false,
custom_ignores: &custom_ignores,
exclude_filenames: false,
language: None,
reranker: "hybrid",
frequency_search: false,
max_results: Some(100),
max_bytes: Some(1_000_000),
max_tokens: Some(100_000),
allow_tests: true,
no_merge: false,
merge_threshold: Some(5),
dry_run: false,
session: None,
timeout: 30,
exact: false,
};
let start_time = Instant::now();
let result = perform_probe(&options);
let duration = start_time.elapsed();
assert!(result.is_ok(), "Search should succeed");
let search_results = result.unwrap();
assert!(
!search_results.results.is_empty(),
"Search should find matches"
);
println!("Parallel search completed in {duration:?} for {file_count} files");
println!("Found {} results", search_results.results.len());
}
#[test]
fn test_structured_patterns_search() {
let temp_dir = tempdir().unwrap();
let base_path = temp_dir.path();
let file_count = 20;
let lines_per_file = 50;
let _file_paths = create_test_files(base_path, file_count, lines_per_file);
let query_plan = probe_code::search::query::create_query_plan(
"search_term_alpha OR search_term_beta",
false,
)
.unwrap();
let patterns = probe_code::search::query::create_structured_patterns(&query_plan);
let custom_ignores: Vec<String> = Vec::new();
let start_time = Instant::now();
let result = search_with_structured_patterns(
base_path,
&query_plan,
&patterns,
&custom_ignores,
true,
None,
);
let duration = start_time.elapsed();
assert!(result.is_ok(), "Structured pattern search should succeed");
let file_term_maps = result.unwrap();
assert!(!file_term_maps.is_empty(), "Search should find matches");
println!("Parallel structured pattern search completed in {duration:?} for {file_count} files");
println!("Found matches in {} files", file_term_maps.len());
}
#[test]
fn test_ast_parallel_processing() {
let temp_dir = tempdir().unwrap();
let base_path = temp_dir.path();
let file_path = base_path.join("large_file.rs");
let mut content = "// Large test file with many top-level nodes\n\n".to_string();
for i in 0..100 {
content.push_str(&format!(
"/// Function documentation for function_{}
fn function_{}() {{
// Function body with search terms
let search_term_alpha = {};
let search_term_beta = {};
println!(\"Processing with search_term_gamma\");
}}
",
i,
i,
i * 2,
i * 3
));
}
fs::write(&file_path, content).unwrap();
let queries = vec![
"search_term_alpha".to_string(),
"search_term_beta".to_string(),
];
let custom_ignores = Vec::new();
let options = SearchOptions {
path: base_path,
queries: &queries,
files_only: false,
custom_ignores: &custom_ignores,
exclude_filenames: false,
language: None,
reranker: "hybrid",
frequency_search: false,
max_results: Some(100),
max_bytes: Some(1_000_000),
max_tokens: Some(100_000),
allow_tests: true,
no_merge: false,
merge_threshold: Some(5),
dry_run: false,
session: None,
timeout: 30,
exact: false,
};
let start_time = Instant::now();
let result = perform_probe(&options);
let duration = start_time.elapsed();
assert!(
result.is_ok(),
"AST parallel processing search should succeed"
);
let search_results = result.unwrap();
assert!(
!search_results.results.is_empty(),
"Search should find matches"
);
println!("AST parallel processing completed in {duration:?}");
println!("Found {} results", search_results.results.len());
}
#[test]
fn test_block_parallel_processing() {
let temp_dir = tempdir().unwrap();
let base_path = temp_dir.path();
let file_path = base_path.join("multi_block.rs");
let mut content = "// Test file with many code blocks\n\n".to_string();
for i in 0..50 {
content.push_str(&format!(
"/// Function with multiple blocks
fn function_with_blocks_{}() {{
// Block 1
{{
let search_term_alpha = {};
println!(\"Block 1\");
}}
// Block 2
{{
let search_term_beta = {};
println!(\"Block 2\");
}}
// Block 3
{{
let search_term_gamma = {};
println!(\"Block 3\");
}}
// Block 4
if true {{
let search_term_delta = {};
println!(\"Block 4\");
}}
}}
",
i,
i,
i * 2,
i * 3,
i * 4
));
}
fs::write(&file_path, content).unwrap();
let queries = vec![
"search_term_alpha".to_string(),
"search_term_beta".to_string(),
];
let custom_ignores = Vec::new();
let options = SearchOptions {
path: base_path,
queries: &queries,
files_only: false,
custom_ignores: &custom_ignores,
exclude_filenames: false,
language: None,
reranker: "hybrid",
frequency_search: false,
max_results: Some(100),
max_bytes: Some(1_000_000),
max_tokens: Some(100_000),
allow_tests: true,
no_merge: false,
merge_threshold: Some(5),
dry_run: false,
session: None,
timeout: 30,
exact: false,
};
let start_time = Instant::now();
let result = perform_probe(&options);
let duration = start_time.elapsed();
assert!(
result.is_ok(),
"Block parallel processing search should succeed"
);
let search_results = result.unwrap();
assert!(
!search_results.results.is_empty(),
"Search should find matches"
);
println!("Block parallel processing completed in {duration:?}");
println!("Found {} results", search_results.results.len());
}