_diffctx/edges/structural/
testing.rs

1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::weights::EDGE_WEIGHTS;
8use crate::types::Fragment;
9
10use super::super::EdgeDict;
11use super::super::base::{self, EdgeBuilder, add_edge_unidirectional, path_to_module};
12
13static IMPORT_RE: Lazy<Regex> =
14    Lazy::new(|| Regex::new(r"(?m)^\s*(?:from\s+([\w.]+)\s+import|import\s+([\w.]+))").unwrap());
15
16fn is_python_test(name: &str) -> bool {
17    name.starts_with("test_") || name.ends_with("_test.py")
18}
19
20fn is_js_test(name: &str, path_str: &str) -> bool {
21    name.contains(".test.") || name.contains(".spec.") || path_str.contains("__tests__")
22}
23
24fn is_rust_test(name: &str, path_str: &str) -> bool {
25    path_str.contains("/tests/") || name == "tests.rs"
26}
27
28fn is_jvm_test(name: &str) -> bool {
29    let stem = if let Some(idx) = name.rfind('.') {
30        &name[..idx]
31    } else {
32        name
33    };
34    let lower = stem.to_lowercase();
35    lower.ends_with("test") || lower.starts_with("test")
36}
37
38fn is_test_file(path: &Path) -> bool {
39    let name = path
40        .file_name()
41        .map(|n| n.to_string_lossy().to_lowercase())
42        .unwrap_or_default();
43    let path_str = path.to_string_lossy().to_lowercase();
44    let ext = base::file_ext(path);
45
46    let lang_match = match ext.as_str() {
47        ".py" => is_python_test(&name),
48        ".js" | ".ts" | ".jsx" | ".tsx" => is_js_test(&name, &path_str),
49        ".rs" => is_rust_test(&name, &path_str),
50        ".java" | ".kt" | ".kts" | ".scala" => is_jvm_test(&name),
51        _ => false,
52    };
53    lang_match || path_str.contains("/tests/") || path_str.contains("/test/")
54}
55
56fn extract_imports(content: &str) -> FxHashSet<String> {
57    let mut imports = FxHashSet::default();
58    for cap in IMPORT_RE.captures_iter(content) {
59        if let Some(m) = cap.get(1) {
60            imports.insert(m.as_str().to_string());
61        }
62        if let Some(m) = cap.get(2) {
63            imports.insert(m.as_str().to_string());
64        }
65    }
66    imports
67}
68
69fn has_direct_import(test_imports: &FxHashSet<String>, src_module: &str) -> bool {
70    if src_module.is_empty() {
71        return false;
72    }
73    let suffix = format!(".{}", src_module);
74    test_imports
75        .iter()
76        .any(|imp| imp == src_module || imp.ends_with(&suffix))
77}
78
79fn extract_target_name_from_test(test_name: &str) -> Option<String> {
80    let lower = test_name.to_lowercase();
81    if lower.starts_with("test_") {
82        return Some(lower[5..].to_string());
83    }
84    if lower.ends_with("_test") {
85        return Some(lower[..lower.len() - 5].to_string());
86    }
87    if lower.contains(".test") {
88        return Some(lower.split(".test").next()?.to_string());
89    }
90    if lower.contains(".spec") {
91        return Some(lower.split(".spec").next()?.to_string());
92    }
93    if test_name.starts_with("Test")
94        && test_name.len() > 4
95        && test_name.as_bytes()[4].is_ascii_uppercase()
96    {
97        return Some(test_name[4..].to_lowercase());
98    }
99    if test_name.ends_with("Tests") && test_name.len() > 5 {
100        return Some(test_name[..test_name.len() - 5].to_lowercase());
101    }
102    if test_name.ends_with("Test") && test_name.len() > 4 {
103        return Some(test_name[..test_name.len() - 4].to_lowercase());
104    }
105    None
106}
107
108pub struct TestEdgeBuilder;
109
110impl EdgeBuilder for TestEdgeBuilder {
111    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
112        let weight_direct = EDGE_WEIGHTS["test_direct"].forward;
113        let weight_naming = EDGE_WEIGHTS["test_naming"].forward;
114        let test_reverse_weight = EDGE_WEIGHTS["test_reverse"].forward;
115
116        let mut test_frags: Vec<&Fragment> = Vec::new();
117        let mut by_base: FxHashMap<String, Vec<&Fragment>> = FxHashMap::default();
118
119        for f in fragments {
120            let path = Path::new(f.path());
121            if is_test_file(path) {
122                test_frags.push(f);
123            } else {
124                let stem = path
125                    .file_stem()
126                    .map(|s| s.to_string_lossy().to_lowercase())
127                    .unwrap_or_default();
128                by_base.entry(stem).or_default().push(f);
129            }
130        }
131
132        let mut module_cache: FxHashMap<String, String> = FxHashMap::default();
133        for src_list in by_base.values() {
134            for sf in src_list {
135                let path_str = sf.path().to_string();
136                module_cache
137                    .entry(path_str)
138                    .or_insert_with_key(|_| path_to_module(Path::new(sf.path()), repo_root));
139            }
140        }
141
142        let mut import_cache: FxHashMap<String, FxHashSet<String>> = FxHashMap::default();
143        for tf in &test_frags {
144            let path_str = tf.path().to_string();
145            import_cache
146                .entry(path_str)
147                .or_insert_with(|| extract_imports(&tf.content));
148        }
149
150        let mut edges: EdgeDict = FxHashMap::default();
151
152        for test_frag in &test_frags {
153            let test_stem = Path::new(test_frag.path())
154                .file_stem()
155                .map(|s| s.to_string_lossy().to_string())
156                .unwrap_or_default();
157            let target_name = match extract_target_name_from_test(&test_stem) {
158                Some(name) => name,
159                None => continue,
160            };
161
162            let test_imports = import_cache
163                .get(test_frag.path())
164                .cloned()
165                .unwrap_or_default();
166
167            for src_frag in by_base.get(&target_name).unwrap_or(&vec![]) {
168                let src_module = module_cache
169                    .get(src_frag.path())
170                    .map(|s| s.as_str())
171                    .unwrap_or("");
172                let weight = if has_direct_import(&test_imports, src_module) {
173                    weight_direct
174                } else {
175                    weight_naming
176                };
177
178                add_edge_unidirectional(&mut edges, &test_frag.id, &src_frag.id, weight);
179                add_edge_unidirectional(
180                    &mut edges,
181                    &src_frag.id,
182                    &test_frag.id,
183                    test_reverse_weight,
184                );
185            }
186        }
187
188        edges
189    }
190
191    fn discover_related_files(
192        &self,
193        changed: &[PathBuf],
194        candidates: &[PathBuf],
195        _repo_root: Option<&Path>,
196        _file_cache: Option<&FxHashMap<PathBuf, String>>,
197    ) -> Vec<PathBuf> {
198        let changed_set: FxHashSet<PathBuf> = changed.iter().cloned().collect();
199        let mut candidate_by_stem: FxHashMap<String, Vec<PathBuf>> = FxHashMap::default();
200        for c in candidates {
201            if !changed_set.contains(c) {
202                let stem = c
203                    .file_stem()
204                    .map(|s| s.to_string_lossy().to_lowercase())
205                    .unwrap_or_default();
206                candidate_by_stem.entry(stem).or_default().push(c.clone());
207            }
208        }
209
210        let mut discovered: Vec<PathBuf> = Vec::new();
211
212        for changed_file in changed {
213            let ext = base::file_ext(changed_file);
214            let stem = changed_file
215                .file_stem()
216                .map(|s| s.to_string_lossy().to_string())
217                .unwrap_or_default();
218
219            if is_test_file(changed_file) {
220                if let Some(target) = extract_target_name_from_test(&stem) {
221                    for c in candidate_by_stem.get(&target).unwrap_or(&vec![]) {
222                        if base::file_ext(c) == ext {
223                            discovered.push(c.clone());
224                        }
225                    }
226                }
227            } else {
228                let stem_lower = stem.to_lowercase();
229                for test_stem in [
230                    format!("test_{}", stem_lower),
231                    format!("{}_test", stem_lower),
232                ] {
233                    for c in candidate_by_stem.get(&test_stem).unwrap_or(&vec![]) {
234                        if base::file_ext(c) == ext && is_test_file(c) {
235                            discovered.push(c.clone());
236                        }
237                    }
238                }
239
240                if matches!(
241                    ext.as_str(),
242                    ".js" | ".ts" | ".jsx" | ".tsx" | ".mjs" | ".cjs"
243                ) {
244                    let stem_test = format!("{stem_lower}.test");
245                    let stem_spec = format!("{stem_lower}.spec");
246                    for c in candidate_by_stem.get(&stem_test).unwrap_or(&vec![]) {
247                        if base::file_ext(c) == ext && is_test_file(c) {
248                            discovered.push(c.clone());
249                        }
250                    }
251                    for c in candidate_by_stem.get(&stem_spec).unwrap_or(&vec![]) {
252                        if base::file_ext(c) == ext && is_test_file(c) {
253                            discovered.push(c.clone());
254                        }
255                    }
256                }
257            }
258        }
259
260        discovered
261    }
262
263    fn category_label(&self) -> Option<&str> {
264        Some("test_edge")
265    }
266}
_diffctx/edges/structural/testing.rs

_diffctx/edges/structural/
testing.rs
