use crate::error::Result;
use crate::graph::GraphStore;
use crate::types::{Edge, Node, NodeId, QueryResult};
use std::collections::VecDeque;
use std::collections::{HashMap, HashSet};
pub struct QueryEngine {
store: Box<dyn GraphStore>,
max_nodes: usize,
}
impl QueryEngine {
pub fn new(store: Box<dyn GraphStore>) -> Self {
QueryEngine {
store,
max_nodes: 10_000,
}
}
pub fn with_max_nodes(store: Box<dyn GraphStore>, max_nodes: usize) -> Self {
QueryEngine { store, max_nodes }
}
pub fn store_ref(&self) -> &dyn GraphStore {
self.store.as_ref()
}
pub fn query_text(
&self,
query: &str,
mode: &str,
depth: usize,
token_budget: Option<usize>,
context_filter: Option<&str>,
) -> Result<QueryResult> {
let seeds = self.store.search(query, 10)?;
let seed_nodes: Vec<Node> = seeds.into_iter().map(|(_, n)| n).collect();
if seed_nodes.is_empty() {
return Ok(QueryResult {
seed_nodes: vec![],
neighborhood: vec![],
edges: vec![],
truncated: false,
});
}
let node_count = self.store.node_count()?;
if node_count > self.max_nodes {
return Ok(QueryResult {
seed_nodes: vec![],
neighborhood: vec![],
edges: vec![],
truncated: true,
});
}
match mode {
"bfs" => self.traverse_bfs(&seed_nodes, depth, context_filter, token_budget),
"dfs" => self.traverse_dfs(&seed_nodes, depth, context_filter, token_budget),
_ => self.traverse_bfs(&seed_nodes, depth, context_filter, token_budget),
}
}
fn traverse_bfs(
&self,
seed_nodes: &[Node],
depth: usize,
context_filter: Option<&str>,
token_budget: Option<usize>,
) -> Result<QueryResult> {
let mut visited: HashSet<NodeId> = HashSet::new();
let mut neighborhood: Vec<Node> = Vec::new();
let mut edges: Vec<Edge> = Vec::new();
for seed in seed_nodes {
visited.insert(seed.id.clone());
neighborhood.push(seed.clone());
}
if depth > 0 {
let mut frontier: Vec<NodeId> = seed_nodes.iter().map(|n| n.id.clone()).collect();
for _ in 0..depth {
let mut next_frontier = Vec::new();
for node_id in &frontier {
if let Ok(neighbors) = self.store.neighbors(node_id, context_filter) {
for (neighbor, edge) in &neighbors {
if visited.insert(neighbor.id.clone()) {
neighborhood.push(neighbor.clone());
next_frontier.push(neighbor.id.clone());
}
edges.push(edge.clone());
}
}
}
frontier = next_frontier;
if frontier.is_empty() {
break;
}
if let Some(budget) = token_budget {
let token_count: usize =
neighborhood.iter().map(|n| n.label.len()).sum::<usize>()
+ edges.iter().map(|e| e.relation.len()).sum::<usize>();
if token_count >= budget {
return Ok(QueryResult {
seed_nodes: seed_nodes.to_vec(),
neighborhood,
edges,
truncated: true,
});
}
}
}
}
Ok(QueryResult {
seed_nodes: seed_nodes.to_vec(),
neighborhood,
edges,
truncated: false,
})
}
fn traverse_dfs(
&self,
seed_nodes: &[Node],
depth: usize,
context_filter: Option<&str>,
token_budget: Option<usize>,
) -> Result<QueryResult> {
let mut visited: HashSet<NodeId> = HashSet::new();
let mut neighborhood: Vec<Node> = Vec::new();
let mut edges: Vec<Edge> = Vec::new();
for seed in seed_nodes {
visited.insert(seed.id.clone());
neighborhood.push(seed.clone());
}
if depth > 0 {
let mut stack: Vec<(NodeId, usize)> =
seed_nodes.iter().map(|n| (n.id.clone(), 0)).collect();
while let Some((node_id, current_depth)) = stack.pop() {
if current_depth >= depth {
continue;
}
if let Ok(neighbors) = self.store.neighbors(&node_id, context_filter) {
for (neighbor, edge) in &neighbors {
if visited.insert(neighbor.id.clone()) {
neighborhood.push(neighbor.clone());
stack.push((neighbor.id.clone(), current_depth + 1));
}
edges.push(edge.clone());
}
}
if let Some(budget) = token_budget {
let token_count: usize =
neighborhood.iter().map(|n| n.label.len()).sum::<usize>()
+ edges.iter().map(|e| e.relation.len()).sum::<usize>();
if token_count >= budget {
return Ok(QueryResult {
seed_nodes: seed_nodes.to_vec(),
neighborhood,
edges,
truncated: true,
});
}
}
}
}
Ok(QueryResult {
seed_nodes: seed_nodes.to_vec(),
neighborhood,
edges,
truncated: false,
})
}
pub fn dijkstra_path(&self, src: &str, tgt: &str) -> Result<Option<(Vec<Node>, f64)>> {
if src == tgt {
return Ok(self.store.get_node(src)?.map(|n| (vec![n], 0.0)));
}
let all_nodes = self.store.get_all_nodes()?;
let all_edges = self.store.get_all_edges()?;
let mut dist: HashMap<&str, f64> = HashMap::new();
let mut prev: HashMap<&str, String> = HashMap::new();
let mut unvisited: HashSet<&str> = HashSet::new();
for node in &all_nodes {
dist.insert(node.id.as_str(), f64::INFINITY);
unvisited.insert(node.id.as_str());
}
dist.insert(src, 0.0);
while !unvisited.is_empty() {
let current = unvisited
.iter()
.min_by(|a, b| {
dist.get(*a)
.unwrap_or(&f64::INFINITY)
.partial_cmp(dist.get(*b).unwrap_or(&f64::INFINITY))
.unwrap_or(std::cmp::Ordering::Equal)
})
.cloned();
let current = match current {
Some(c) => c,
None => break,
};
if current == tgt {
break;
}
unvisited.remove(current);
let current_dist = *dist.get(current).unwrap_or(&f64::INFINITY);
if current_dist == f64::INFINITY {
break;
}
for edge in &all_edges {
let neighbor = if edge.source == current {
Some(edge.target.as_str())
} else if edge.target == current {
Some(edge.source.as_str())
} else {
None
};
if let Some(neighbor) = neighbor {
if unvisited.contains(neighbor) {
let alt = current_dist + edge.weight;
if alt < *dist.get(neighbor).unwrap_or(&f64::INFINITY) {
dist.insert(neighbor, alt);
prev.insert(neighbor, current.to_string());
}
}
}
}
}
if prev.contains_key(tgt) || src == tgt {
let mut path_nodes = Vec::new();
let mut current = tgt.to_string();
while current != src {
if let Some(node) = self.store.get_node(¤t)? {
path_nodes.push(node);
}
current = prev.get(current.as_str()).cloned().unwrap_or_default();
}
if let Some(node) = self.store.get_node(src)? {
path_nodes.push(node);
}
path_nodes.reverse();
let total_weight = *dist.get(tgt).unwrap_or(&f64::INFINITY);
Ok(Some((path_nodes, total_weight)))
} else {
Ok(None)
}
}
pub fn dfs(&self, src: &str, depth: usize, relation_filter: Option<&str>) -> Result<Vec<Node>> {
let mut visited: HashSet<NodeId> = HashSet::new();
let mut result: Vec<Node> = Vec::new();
let mut stack: Vec<(NodeId, usize)> = vec![(src.to_string(), 0)];
while let Some((node_id, d)) = stack.pop() {
if d > depth {
continue;
}
if !visited.insert(node_id.clone()) {
continue;
}
if let Some(node) = self.store.get_node(&node_id)? {
result.push(node);
}
if d < depth {
if let Ok(neighbors) = self.store.neighbors(&node_id, relation_filter) {
for (neighbor, _) in &neighbors {
if !visited.contains(&neighbor.id) {
stack.push((neighbor.id.clone(), d + 1));
}
}
}
}
}
Ok(result)
}
pub fn get_node_by_label(&self, label: &str) -> Result<Option<Node>> {
let nodes = self.store.search(label, 1)?;
Ok(nodes.into_iter().next().map(|(_, n)| n))
}
pub fn get_node_by_id(&self, id: &str) -> Result<Option<Node>> {
self.store.get_node(id)
}
pub fn resolve_node(&self, id_or_label: &str) -> Result<Option<Node>> {
if let Some(node) = self.store.get_node(id_or_label)? {
return Ok(Some(node));
}
self.get_node_by_label(id_or_label)
}
pub fn get_neighbors(
&self,
id: &str,
relation_filter: Option<&str>,
) -> Result<Vec<(Node, Edge)>> {
self.store.neighbors(id, relation_filter)
}
pub fn shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
self.store.shortest_path(src, tgt)
}
pub fn graph_stats(&self) -> Result<String> {
let node_count = self.store.node_count()?;
let edge_count = self.store.edge_count()?;
Ok(format!("Nodes: {}, Edges: {}", node_count, edge_count))
}
pub fn get_community(&self, community_id: usize) -> Result<Vec<Node>> {
let nodes = self.store.get_all_nodes()?;
Ok(nodes
.into_iter()
.filter(|n| n.community == Some(community_id))
.collect())
}
pub fn god_nodes(&self, top_n: usize) -> Result<Vec<Node>> {
let edges = self.store.get_all_edges()?;
let all_nodes = self.store.get_all_nodes()?;
let mut degree: HashMap<NodeId, usize> = HashMap::new();
for edge in &edges {
*degree.entry(edge.source.clone()).or_insert(0) += 1;
*degree.entry(edge.target.clone()).or_insert(0) += 1;
}
let mut ranked: Vec<(usize, NodeId)> = degree.into_iter().map(|(k, v)| (v, k)).collect();
ranked.sort_by_key(|k| std::cmp::Reverse(k.0));
ranked.truncate(top_n);
let id_set: HashSet<&str> = ranked.iter().map(|(_, id)| id.as_str()).collect();
let mut result: Vec<Node> = all_nodes
.into_iter()
.filter(|n| id_set.contains(n.id.as_str()))
.collect();
result.sort_by(|a, b| {
let da = ranked
.iter()
.find(|(_, id)| id == &a.id)
.map(|(d, _)| *d)
.unwrap_or(0);
let db = ranked
.iter()
.find(|(_, id)| id == &b.id)
.map(|(d, _)| *d)
.unwrap_or(0);
db.cmp(&da)
});
Ok(result)
}
pub fn shortest_path_with_max_hops(
&self,
src: &str,
tgt: &str,
max_hops: usize,
) -> Result<Option<Vec<Node>>> {
if src == tgt {
return self.store.get_node(src).map(|n| n.map(|n| vec![n]));
}
let mut visited: HashSet<NodeId> = HashSet::new();
let mut queue: VecDeque<(NodeId, usize)> = VecDeque::new();
let mut parent: HashMap<NodeId, NodeId> = HashMap::new();
visited.insert(src.to_string());
queue.push_back((src.to_string(), 0));
while let Some((current, depth)) = queue.pop_front() {
if depth >= max_hops {
continue;
}
let neighbors = self.store.neighbors(¤t, None)?;
for (neighbor, _) in &neighbors {
if visited.insert(neighbor.id.clone()) {
parent.insert(neighbor.id.clone(), current.clone());
if neighbor.id == tgt {
let mut path = Vec::new();
let mut cur = tgt.to_string();
while cur != src {
if let Some(n) = self.store.get_node(&cur)? {
path.push(n);
}
cur = parent.get(&cur).cloned().unwrap_or_default();
}
if let Some(n) = self.store.get_node(src)? {
path.push(n);
}
path.reverse();
return Ok(Some(path));
}
queue.push_back((neighbor.id.clone(), depth + 1));
}
}
}
Ok(None)
}
pub fn infer_context_filter(question: &str) -> Option<String> {
const CONTEXT_HINTS: &[(&str, &[&str])] = &[
(
"call",
&["call", "calls", "called", "invoke", "invokes", "invoked"],
),
(
"import",
&["import", "imports", "imported", "module", "modules"],
),
(
"field",
&[
"field",
"fields",
"member",
"members",
"property",
"properties",
],
),
(
"parameter_type",
&[
"parameter",
"parameters",
"param",
"params",
"argument",
"arguments",
],
),
("return_type", &["return", "returns", "returned"]),
(
"generic_arg",
&["generic", "generics", "template", "templates"],
),
];
let tokens: std::collections::HashSet<String> = question
.split(|c: char| !c.is_alphanumeric())
.filter(|s| !s.is_empty())
.map(|s| s.to_lowercase())
.collect();
for (context, hints) in CONTEXT_HINTS {
if hints.iter().any(|h| tokens.contains(*h)) {
return Some(context.to_string());
}
}
None
}
pub fn resolve_context_filter(question: &str, explicit: Option<&str>) -> Option<String> {
const ALIASES: &[(&str, &str)] = &[
("param", "parameter_type"),
("params", "parameter_type"),
("parameter", "parameter_type"),
("parameters", "parameter_type"),
("argument", "parameter_type"),
("arguments", "parameter_type"),
("arg", "parameter_type"),
("args", "parameter_type"),
("return", "return_type"),
("returns", "return_type"),
("returned", "return_type"),
("generic", "generic_arg"),
("generics", "generic_arg"),
("template", "generic_arg"),
("templates", "generic_arg"),
("calls", "call"),
("called", "call"),
("invoke", "call"),
("invocation", "call"),
("fields", "field"),
("property", "field"),
("properties", "field"),
("member", "field"),
("members", "field"),
("imports", "import"),
("imported", "import"),
("module", "import"),
("modules", "import"),
("exports", "export"),
("exported", "export"),
];
if let Some(exp) = explicit {
let key = exp.trim().to_lowercase();
let normalized = ALIASES
.iter()
.find(|(k, _)| *k == key)
.map(|(_, v)| v.to_string())
.unwrap_or(key);
return Some(normalized);
}
Self::infer_context_filter(question)
}
pub fn detailed_stats(&self) -> Result<String> {
let nodes = self.store.get_all_nodes()?;
let edges = self.store.get_all_edges()?;
let mut file_types: HashMap<&str, usize> = HashMap::new();
for node in &nodes {
let ext = node.source_file.rsplit('.').next().unwrap_or("unknown");
*file_types.entry(ext).or_insert(0) += 1;
}
let mut rel_types: HashMap<&str, usize> = HashMap::new();
for edge in &edges {
*rel_types.entry(edge.relation.as_str()).or_insert(0) += 1;
}
let mut stats = String::new();
stats.push_str(&format!("Nodes: {}, Edges: {}\n", nodes.len(), edges.len()));
stats.push_str(&format!("File types: {:?}\n", file_types));
stats.push_str(&format!("Relation types: {:?}\n", rel_types));
Ok(stats)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::Node;
use std::collections::HashMap;
use std::sync::Mutex;
struct MockStore {
nodes: Mutex<HashMap<NodeId, Node>>,
edges: Mutex<Vec<Edge>>,
}
impl MockStore {
fn new() -> Self {
MockStore {
nodes: Mutex::new(HashMap::new()),
edges: Mutex::new(Vec::new()),
}
}
}
impl GraphStore for MockStore {
fn dijkstra_shortest_path(&self, _src: &str, _tgt: &str) -> Result<Option<Vec<Node>>> {
Ok(None)
}
fn add_node(&self, node: Node) -> Result<()> {
self.nodes.lock().unwrap().insert(node.id.clone(), node);
Ok(())
}
fn add_edge(&self, edge: Edge) -> Result<()> {
self.edges.lock().unwrap().push(edge);
Ok(())
}
fn get_node(&self, id: &str) -> Result<Option<Node>> {
Ok(self.nodes.lock().unwrap().get(id).cloned())
}
fn get_all_nodes(&self) -> Result<Vec<Node>> {
Ok(self.nodes.lock().unwrap().values().cloned().collect())
}
fn get_all_edges(&self) -> Result<Vec<Edge>> {
Ok(self.edges.lock().unwrap().clone())
}
fn neighbors(&self, id: &str, filter: Option<&str>) -> Result<Vec<(Node, Edge)>> {
let edges = self.edges.lock().unwrap();
let nodes = self.nodes.lock().unwrap();
let mut result = Vec::new();
for edge in edges.iter() {
if let Some(f) = filter {
if edge.relation != f {
continue;
}
}
if edge.source == id {
if let Some(node) = nodes.get(&edge.target) {
result.push((node.clone(), edge.clone()));
}
} else if edge.target == id {
if let Some(node) = nodes.get(&edge.source) {
result.push((node.clone(), edge.clone()));
}
}
}
Ok(result)
}
fn search(&self, query: &str, _top_k: usize) -> Result<Vec<(f64, Node)>> {
let nodes = self.nodes.lock().unwrap();
let q = query.to_lowercase();
let mut results = Vec::new();
for node in nodes.values() {
if node.label.to_lowercase().contains(&q) || node.id.to_lowercase().contains(&q) {
results.push((1.0, node.clone()));
}
}
Ok(results)
}
fn shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
if src == tgt {
return Ok(self
.nodes
.lock()
.unwrap()
.get(src)
.cloned()
.map(|n| vec![n]));
}
let nodes = self.nodes.lock().unwrap();
let edges = self.edges.lock().unwrap();
let mut visited = HashSet::new();
let mut queue = VecDeque::new();
let mut parent: HashMap<String, String> = HashMap::new();
visited.insert(src.to_string());
queue.push_back(src.to_string());
while let Some(current) = queue.pop_front() {
if current == tgt {
let mut path = Vec::new();
let mut cur = tgt.to_string();
while cur != src {
if let Some(n) = nodes.get(&cur) {
path.push(n.clone());
}
cur = parent.get(&cur).cloned().unwrap_or_default();
}
if let Some(n) = nodes.get(src) {
path.push(n.clone());
}
path.reverse();
return Ok(Some(path));
}
for edge in edges.iter() {
let neighbor = if edge.source == current {
Some(edge.target.clone())
} else if edge.target == current {
Some(edge.source.clone())
} else {
None
};
if let Some(neighbor) = neighbor {
if visited.insert(neighbor.clone()) {
parent.insert(neighbor.clone(), current.clone());
queue.push_back(neighbor);
}
}
}
}
Ok(None)
}
fn subgraph(&self, _node_ids: &[&str]) -> Result<(Vec<Node>, Vec<Edge>)> {
Ok((vec![], vec![]))
}
fn node_count(&self) -> Result<usize> {
Ok(self.nodes.lock().unwrap().len())
}
fn edge_count(&self) -> Result<usize> {
Ok(self.edges.lock().unwrap().len())
}
fn remove_node(&self, id: &str) -> Result<()> {
self.nodes.lock().unwrap().remove(id);
Ok(())
}
fn remove_edge(&self, source: &str, target: &str, relation: &str) -> Result<()> {
let mut edges = self.edges.lock().unwrap();
edges.retain(|e| !(e.source == source && e.target == target && e.relation == relation));
Ok(())
}
fn clear(&self) -> Result<()> {
self.nodes.lock().unwrap().clear();
self.edges.lock().unwrap().clear();
Ok(())
}
}
fn make_node(id: &str, label: &str) -> Node {
Node {
id: id.to_string(),
label: label.to_string(),
file_type: "code".to_string(),
source_file: "test.py".to_string(),
source_location: None,
community: None,
rationale: None,
docstring: None,
metadata: HashMap::new(),
}
}
fn make_edge(src: &str, tgt: &str, relation: &str, weight: f64) -> Edge {
Edge {
source: src.to_string(),
target: tgt.to_string(),
relation: relation.to_string(),
confidence: "EXTRACTED".to_string(),
source_file: Some("test.py".to_string()),
weight,
context: None,
}
}
#[test]
fn test_get_node_by_id() {
let store = MockStore::new();
store.add_node(make_node("auth", "AuthService")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let node = engine.get_node_by_id("auth").unwrap();
assert!(node.is_some());
assert_eq!(node.unwrap().label, "AuthService");
}
#[test]
fn test_get_node_not_found() {
let store = MockStore::new();
let engine = QueryEngine::new(Box::new(store));
let node = engine.get_node_by_id("nonexistent").unwrap();
assert!(node.is_none());
}
#[test]
fn test_graph_stats() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let stats = engine.graph_stats().unwrap();
assert!(stats.contains("Nodes: 2"));
}
#[test]
fn test_detailed_stats() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let stats = engine.detailed_stats().unwrap();
assert!(stats.contains("Nodes: 2"));
assert!(stats.contains("Edges: 1"));
}
#[test]
fn test_bfs_traversal() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("a", "bfs", 2, None, None).unwrap();
assert_eq!(result.seed_nodes.len(), 1);
assert_eq!(result.seed_nodes[0].id, "a");
assert!(result.neighborhood.len() >= 2);
}
#[test]
fn test_dfs_traversal() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("a", "dfs", 2, None, None).unwrap();
assert_eq!(result.seed_nodes.len(), 1);
assert!(result.neighborhood.len() >= 2);
}
#[test]
fn test_dfs_method() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let nodes = engine.dfs("a", 3, None).unwrap();
assert_eq!(nodes.len(), 3);
}
#[test]
fn test_dijkstra_path() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.5))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 2.5))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.dijkstra_path("a", "c").unwrap();
assert!(result.is_some());
let (path, weight) = result.unwrap();
assert_eq!(path.len(), 3);
assert!((weight - 4.0).abs() < 0.01);
}
#[test]
fn test_dijkstra_path_same_node() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.dijkstra_path("a", "a").unwrap();
assert!(result.is_some());
let (path, weight) = result.unwrap();
assert_eq!(path.len(), 1);
assert!((weight - 0.0).abs() < 0.01);
}
#[test]
fn test_dijkstra_path_no_path() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.dijkstra_path("a", "b").unwrap();
assert!(result.is_none());
}
#[test]
fn test_context_filter() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store.add_edge(make_edge("a", "b", "imports", 1.0)).unwrap();
store.add_edge(make_edge("a", "c", "calls", 1.0)).unwrap();
let engine = QueryEngine::new(Box::new(store));
let neighbors = engine.get_neighbors("a", Some("imports")).unwrap();
assert_eq!(neighbors.len(), 1);
assert_eq!(neighbors[0].0.label, "B");
}
#[test]
fn test_token_budget() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("a", "bfs", 1, Some(5), None).unwrap();
assert_eq!(result.seed_nodes.len(), 1);
}
#[test]
fn test_query_exact_match() {
let store = MockStore::new();
store.add_node(make_node("a1", "AuthService")).unwrap();
store.add_node(make_node("u1", "UserService")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine
.query_text("AuthService", "bfs", 0, None, None)
.unwrap();
assert_eq!(result.seed_nodes.len(), 1);
assert_eq!(result.seed_nodes[0].label, "AuthService");
}
#[test]
fn test_query_prefix_match() {
let store = MockStore::new();
store.add_node(make_node("a1", "AuthService")).unwrap();
store.add_node(make_node("a2", "AuthMiddleware")).unwrap();
store.add_node(make_node("u1", "UserService")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("Auth", "bfs", 0, None, None).unwrap();
assert_eq!(result.seed_nodes.len(), 2);
for n in &result.seed_nodes {
assert!(n.label.starts_with("Auth"));
}
}
#[test]
fn test_query_substring_match() {
let store = MockStore::new();
store.add_node(make_node("a1", "AuthService")).unwrap();
store.add_node(make_node("u1", "UserService")).unwrap();
store.add_node(make_node("r1", "RateLimiter")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("Service", "bfs", 0, None, None).unwrap();
assert_eq!(result.seed_nodes.len(), 2);
for n in &result.seed_nodes {
assert!(n.label.contains("Service"));
}
}
#[test]
fn test_query_bfs_depth_1() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("a", "bfs", 1, None, None).unwrap();
assert_eq!(
result.neighborhood.len(),
2,
"depth 1: seed + immediate neighbor"
);
assert!(result.neighborhood.iter().any(|n| n.id == "b"));
assert!(
!result.neighborhood.iter().any(|n| n.id == "c"),
"C should not be reachable at depth 1"
);
}
#[test]
fn test_query_bfs_depth_3() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine.query_text("a", "bfs", 3, None, None).unwrap();
assert_eq!(result.neighborhood.len(), 3, "depth 3: all nodes reachable");
assert!(
result.neighborhood.iter().any(|n| n.id == "c"),
"C should be reachable at depth 3"
);
}
#[test]
fn test_query_context_filter_call() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store.add_edge(make_edge("a", "b", "imports", 1.0)).unwrap();
store.add_edge(make_edge("a", "c", "calls", 1.0)).unwrap();
let engine = QueryEngine::new(Box::new(store));
let result = engine
.query_text("a", "bfs", 1, None, Some("calls"))
.unwrap();
assert_eq!(result.seed_nodes.len(), 1);
let neighbor_ids: Vec<&str> = result.neighborhood.iter().map(|n| n.id.as_str()).collect();
assert!(
neighbor_ids.contains(&"c"),
"call target should be in neighborhood"
);
assert!(
!neighbor_ids.contains(&"b"),
"import target should be filtered out"
);
}
#[test]
fn test_query_context_filter_inferred() {
assert_eq!(
QueryEngine::infer_context_filter("who calls AuthService"),
Some("call".to_string())
);
assert_eq!(
QueryEngine::infer_context_filter("what imports UserModule"),
Some("import".to_string())
);
assert_eq!(
QueryEngine::infer_context_filter("list fields of Foo"),
Some("field".to_string())
);
assert_eq!(
QueryEngine::infer_context_filter("show me AuthService"),
None
);
}
#[test]
fn test_query_context_filter_explicit() {
assert_eq!(
QueryEngine::resolve_context_filter("who calls AuthService", Some("imports")),
Some("import".to_string())
);
assert_eq!(
QueryEngine::resolve_context_filter("who calls AuthService", None),
Some("call".to_string())
);
assert_eq!(
QueryEngine::resolve_context_filter("show me AuthService", None),
None
);
}
#[test]
fn test_get_node_by_label() {
let store = MockStore::new();
store.add_node(make_node("a1", "AuthService")).unwrap();
let engine = QueryEngine::new(Box::new(store));
let node = engine.get_node_by_label("AuthService").unwrap();
assert!(node.is_some());
assert_eq!(node.unwrap().label, "AuthService");
}
#[test]
fn test_get_neighbors_basic() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store.add_edge(make_edge("a", "c", "calls", 1.0)).unwrap();
let engine = QueryEngine::new(Box::new(store));
let neighbors = engine.get_neighbors("a", None).unwrap();
assert_eq!(neighbors.len(), 2);
}
#[test]
fn test_get_community() {
let store = MockStore::new();
let mut n1 = make_node("a", "A");
n1.community = Some(1);
let mut n2 = make_node("b", "B");
n2.community = Some(1);
let n3 = make_node("c", "C");
store.add_node(n1).unwrap();
store.add_node(n2).unwrap();
store.add_node(n3).unwrap();
let engine = QueryEngine::new(Box::new(store));
let members = engine.get_community(1).unwrap();
assert_eq!(members.len(), 2);
for n in &members {
assert_eq!(n.community, Some(1));
}
}
#[test]
fn test_god_nodes() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("a", "c", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let gods = engine.god_nodes(2).unwrap();
assert_eq!(gods.len(), 2);
}
#[test]
fn test_shortest_path_found() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let path = engine.shortest_path("a", "c").unwrap();
assert!(path.is_some());
assert_eq!(path.unwrap().len(), 3);
}
#[test]
fn test_shortest_path_max_hops() {
let store = MockStore::new();
store.add_node(make_node("a", "A")).unwrap();
store.add_node(make_node("b", "B")).unwrap();
store.add_node(make_node("c", "C")).unwrap();
store.add_node(make_node("d", "D")).unwrap();
store
.add_edge(make_edge("a", "b", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("b", "c", "connects", 1.0))
.unwrap();
store
.add_edge(make_edge("c", "d", "connects", 1.0))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let limited = engine.shortest_path_with_max_hops("a", "d", 2).unwrap();
assert!(limited.is_none(), "path exceeds max_hops");
let found = engine.shortest_path_with_max_hops("a", "d", 3).unwrap();
assert!(found.is_some());
assert_eq!(found.unwrap().len(), 4);
}
#[test]
fn test_resolve_node_by_label_when_id_missing() {
let store = MockStore::new();
store
.add_node(make_node("api_authservice", "AuthService"))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let node = engine.resolve_node("AuthService").unwrap();
assert!(node.is_some(), "resolve_node must find node by label");
assert_eq!(node.unwrap().id, "api_authservice");
}
#[test]
fn test_resolve_node_by_exact_id() {
let store = MockStore::new();
store
.add_node(make_node("api_authservice", "AuthService"))
.unwrap();
let engine = QueryEngine::new(Box::new(store));
let node = engine.resolve_node("api_authservice").unwrap();
assert!(node.is_some());
assert_eq!(node.unwrap().id, "api_authservice");
}
#[test]
fn test_resolve_node_not_found() {
let store = MockStore::new();
let engine = QueryEngine::new(Box::new(store));
let node = engine.resolve_node("ghost").unwrap();
assert!(node.is_none());
}
}