Skip to main content

codemem_engine/enrichment/
test_mapping.rs

1//! Test-to-code mapping: link tests to tested symbols, identify untested public functions.
2
3use super::EnrichResult;
4use crate::CodememEngine;
5use codemem_core::{CodememError, Edge, GraphBackend, GraphNode, NodeKind, RelationshipType};
6use serde_json::json;
7use std::collections::{HashMap, HashSet};
8
9impl CodememEngine {
10    /// Map test functions to the code they test and identify untested public functions.
11    ///
12    /// For Test-kind nodes, infers tested symbols by naming convention (`test_foo` -> `foo`)
13    /// and by CALLS edges. Creates RELATES_TO edges between test and tested symbols.
14    /// Produces Insight memories for files with untested public functions.
15    pub fn enrich_test_mapping(
16        &self,
17        namespace: Option<&str>,
18    ) -> Result<EnrichResult, CodememError> {
19        let all_nodes;
20        let mut test_edges_info: Vec<(String, String)> = Vec::new();
21
22        {
23            let graph = self.lock_graph()?;
24            all_nodes = graph.get_all_nodes();
25
26            // Collect test nodes and non-test function/method nodes
27            let test_nodes: Vec<&GraphNode> = all_nodes
28                .iter()
29                .filter(|n| n.kind == NodeKind::Test)
30                .collect();
31            // Index by simple name (last segment of qualified name)
32            let mut fn_by_simple_name: HashMap<String, Vec<&GraphNode>> = HashMap::new();
33            for node in all_nodes
34                .iter()
35                .filter(|n| matches!(n.kind, NodeKind::Function | NodeKind::Method))
36            {
37                let simple = node
38                    .label
39                    .rsplit("::")
40                    .next()
41                    .unwrap_or(&node.label)
42                    .to_string();
43                fn_by_simple_name.entry(simple).or_default().push(node);
44            }
45
46            for test_node in &test_nodes {
47                // Extract what this test might be testing from its name
48                let test_name = test_node
49                    .label
50                    .rsplit("::")
51                    .next()
52                    .unwrap_or(&test_node.label);
53
54                // Convention: test_foo tests foo, test_foo_bar tests foo_bar
55                let tested_name = test_name
56                    .strip_prefix("test_")
57                    .or_else(|| test_name.strip_prefix("test"))
58                    .unwrap_or("");
59
60                if !tested_name.is_empty() {
61                    // Check by simple name
62                    if let Some(targets) = fn_by_simple_name.get(tested_name) {
63                        for target in targets {
64                            test_edges_info.push((test_node.id.clone(), target.id.clone()));
65                        }
66                    }
67                }
68
69                // Also check CALLS edges from the test to find tested symbols
70                if let Ok(edges) = graph.get_edges(&test_node.id) {
71                    for edge in &edges {
72                        if edge.relationship == RelationshipType::Calls && edge.src == test_node.id
73                        {
74                            // Only link to function/method nodes
75                            if let Ok(Some(dst_node)) = graph.get_node(&edge.dst) {
76                                if matches!(dst_node.kind, NodeKind::Function | NodeKind::Method) {
77                                    test_edges_info
78                                        .push((test_node.id.clone(), dst_node.id.clone()));
79                                }
80                            }
81                        }
82                    }
83                }
84            }
85        }
86
87        // Dedup edges
88        let unique_edges: HashSet<(String, String)> = test_edges_info.into_iter().collect();
89
90        // Create RELATES_TO edges for test mappings
91        let mut edges_created = 0;
92        {
93            let mut graph = self.lock_graph()?;
94            let now = chrono::Utc::now();
95            for (test_id, target_id) in &unique_edges {
96                let edge_id = format!("test-map:{test_id}->{target_id}");
97                // Skip if edge already exists
98                if graph.get_node(test_id).ok().flatten().is_none()
99                    || graph.get_node(target_id).ok().flatten().is_none()
100                {
101                    continue;
102                }
103                let edge = Edge {
104                    id: edge_id,
105                    src: test_id.clone(),
106                    dst: target_id.clone(),
107                    relationship: RelationshipType::RelatesTo,
108                    weight: 0.8,
109                    properties: HashMap::from([("test_mapping".into(), json!(true))]),
110                    created_at: now,
111                    valid_from: None,
112                    valid_to: None,
113                };
114                let _ = self.storage.insert_graph_edge(&edge);
115                if graph.add_edge(edge).is_ok() {
116                    edges_created += 1;
117                }
118            }
119        }
120
121        // Identify untested public functions per file
122        let tested_ids: HashSet<String> = unique_edges.iter().map(|(_, t)| t.clone()).collect();
123        let mut untested_by_file: HashMap<String, Vec<String>> = HashMap::new();
124
125        for node in &all_nodes {
126            if !matches!(node.kind, NodeKind::Function | NodeKind::Method) {
127                continue;
128            }
129            let visibility = node
130                .payload
131                .get("visibility")
132                .and_then(|v| v.as_str())
133                .unwrap_or("private");
134            if visibility != "public" {
135                continue;
136            }
137            if tested_ids.contains(&node.id) {
138                continue;
139            }
140            let file_path = node
141                .payload
142                .get("file_path")
143                .and_then(|v| v.as_str())
144                .unwrap_or("unknown")
145                .to_string();
146            untested_by_file
147                .entry(file_path)
148                .or_default()
149                .push(node.label.clone());
150        }
151
152        let mut insights_stored = 0;
153        for (file_path, untested) in &untested_by_file {
154            if untested.is_empty() {
155                continue;
156            }
157            let names: Vec<&str> = untested.iter().take(10).map(|s| s.as_str()).collect();
158            let suffix = if untested.len() > 10 {
159                format!(" (and {} more)", untested.len() - 10)
160            } else {
161                String::new()
162            };
163            let content = format!(
164                "Untested public functions in {}: {}{}",
165                file_path,
166                names.join(", "),
167                suffix
168            );
169            if self
170                .store_insight(
171                    &content,
172                    "testing",
173                    &[],
174                    0.6,
175                    namespace,
176                    &[format!("file:{file_path}")],
177                )
178                .is_some()
179            {
180                insights_stored += 1;
181            }
182        }
183
184        self.save_index();
185
186        Ok(EnrichResult {
187            insights_stored,
188            details: json!({
189                "test_edges_created": edges_created,
190                "files_with_untested": untested_by_file.len(),
191                "insights_stored": insights_stored,
192            }),
193        })
194    }
195}