tests: Added tests for graph-based RAG

src/rag/graph.rs

@@ -732,4 +732,143 @@ mod tests {
         assert!(docs.contains(&DocumentId(0)));
         assert!(docs.contains(&DocumentId(1)));
     }
+
+    #[test]
+    fn compact_preserves_edges_between_survivors() {
+        let mut kg = KnowledgeGraph::default();
+        kg.merge(doc(0), extraction(vec![entity("A", "CONCEPT")], vec![]));
+        kg.merge(
+            doc(1),
+            extraction(
+                vec![entity("B", "CONCEPT"), entity("C", "CONCEPT")],
+                vec![rel("B", "C", "linked", 0.8)],
+            ),
+        );
+        kg.remove_documents(&[doc(0)]);
+        let b_raw = kg.entity_index["b"];
+        let c_raw = kg.entity_index["c"];
+        let b_idx = NodeIndex::new(b_raw as usize);
+        let c_idx = NodeIndex::new(c_raw as usize);
+        assert_eq!(
+            kg.graph.edges_connecting(b_idx, c_idx).count(),
+            1,
+            "B→C edge should survive compaction"
+        );
+    }
+
+    #[test]
+    fn expand_two_hops_reaches_transitive_neighbor() {
+        let mut kg = KnowledgeGraph::default();
+        kg.merge(
+            doc(0),
+            extraction(
+                vec![
+                    entity("A", "CONCEPT"),
+                    entity("B", "CONCEPT"),
+                    entity("C", "CONCEPT"),
+                ],
+                vec![rel("A", "B", "uses", 1.0), rel("B", "C", "uses", 0.5)],
+            ),
+        );
+        let a_raw = kg.entity_index["a"];
+        let c_raw = kg.entity_index["c"];
+
+        let one_hop = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1);
+        assert!(
+            !one_hop.contains_key(&c_raw),
+            "C should not be reachable at 1 hop"
+        );
+
+        let two_hop = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 2);
+        assert!(
+            two_hop.contains_key(&c_raw),
+            "C should be reachable at 2 hops"
+        );
+        let c_score = two_hop[&c_raw];
+        assert!(
+            (c_score - 0.5).abs() < 1e-6,
+            "C score should be 1.0 * 1.0 * 0.5 = 0.5, got {c_score}"
+        );
+    }
+
+    #[test]
+    fn merge_clamps_edge_weight_above_one() {
+        let mut kg = KnowledgeGraph::default();
+        kg.merge(
+            doc(0),
+            extraction(
+                vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
+                vec![rel("A", "B", "uses", 1.5)],
+            ),
+        );
+        let a_raw = kg.entity_index["a"];
+        let b_raw = kg.entity_index["b"];
+        let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1);
+        let b_score = result[&b_raw];
+        assert!(
+            (b_score - 1.0).abs() < 1e-6,
+            "weight 1.5 clamped to 1.0: b_score should be 1.0, got {b_score}"
+        );
+    }
+
+    #[test]
+    fn merge_clamps_edge_weight_below_zero() {
+        let mut kg = KnowledgeGraph::default();
+        kg.merge(
+            doc(0),
+            extraction(
+                vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
+                vec![rel("A", "B", "uses", -0.5)],
+            ),
+        );
+        let a_raw = kg.entity_index["a"];
+        let b_raw = kg.entity_index["b"];
+        let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1);
+        let b_score = result.get(&b_raw).copied().unwrap_or(0.0);
+        assert!(
+            b_score.abs() < 1e-6,
+            "weight -0.5 clamped to 0.0: b_score should be 0.0, got {b_score}"
+        );
+    }
+
+    #[test]
+    fn merge_fills_missing_description_from_later_chunk() {
+        let mut kg = KnowledgeGraph::default();
+        kg.merge(
+            doc(0),
+            extraction(vec![entity("Python", "TECHNOLOGY")], vec![]),
+        );
+        kg.merge(
+            doc(1),
+            ExtractionResult {
+                entities: vec![ExtractedEntity {
+                    name: "python".to_string(),
+                    entity_type: "TECHNOLOGY".to_string(),
+                    description: Some("A general-purpose language".to_string()),
+                }],
+                relationships: vec![],
+            },
+        );
+        let raw = kg.entity_index["python"];
+        let desc = &kg.graph[NodeIndex::new(raw as usize)].description;
+        assert_eq!(
+            desc.as_deref(),
+            Some("A general-purpose language"),
+            "description should be backfilled from later chunk"
+        );
+    }
+
+    #[test]
+    fn remove_all_documents_empties_graph() {
+        let mut kg = KnowledgeGraph::default();
+        kg.merge(doc(0), extraction(vec![entity("A", "CONCEPT")], vec![]));
+        kg.merge(doc(1), extraction(vec![entity("B", "CONCEPT")], vec![]));
+        kg.remove_documents(&[doc(0), doc(1)]);
+        assert_eq!(kg.graph.node_count(), 0, "all nodes should be removed");
+        assert_eq!(kg.entity_index.len(), 0, "entity index should be empty");
+        assert!(
+            kg.document_entities.is_empty(),
+            "document_entities should be empty"
+        );
+    }
 }

src/rag/mod.rs

@@ -1781,4 +1781,91 @@ mod tests {
         assert_eq!(file_idx, 0);
         assert_eq!(doc_idx, 0);
     }
+
+    #[test]
+    fn rag_data_del_removes_graph_entities() {
+        use super::graph::{ExtractedEntity, ExtractionResult};
+        let mut data = RagData::new(
+            "m".into(),
+            100,
+            10,
+            None,
+            5,
+            None,
+            GraphRagConfig::default(),
+        );
+        let file = RagFile {
+            hash: "abc".into(),
+            path: "test.txt".into(),
+            documents: vec![RagDocument::new("Python is great")],
+        };
+        data.files.insert(0, file);
+        let doc_id = DocumentId::new(0, 0);
+        data.knowledge_graph.merge(
+            doc_id,
+            ExtractionResult {
+                entities: vec![ExtractedEntity {
+                    name: "Python".to_string(),
+                    entity_type: "TECHNOLOGY".to_string(),
+                    description: None,
+                }],
+                relationships: vec![],
+            },
+        );
+        assert!(
+            data.knowledge_graph.entity_index.contains_key("python"),
+            "entity should exist before del"
+        );
+        data.del(vec![0]);
+        assert!(
+            !data.knowledge_graph.entity_index.contains_key("python"),
+            "entity should be removed after del"
+        );
+    }
+
+    #[test]
+    fn reciprocal_rank_fusion_empty_lists() {
+        let result = super::reciprocal_rank_fusion(vec![], vec![], 5);
+        assert!(result.is_empty(), "empty input should produce empty output");
+    }
+
+    #[test]
+    fn reciprocal_rank_fusion_deduplicates_across_signals() {
+        let doc_a = DocumentId::new(0, 0);
+        let doc_b = DocumentId::new(0, 1);
+        let result = super::reciprocal_rank_fusion(
+            vec![vec![doc_a, doc_b], vec![doc_a, doc_b]],
+            vec![1.0, 1.0],
+            5,
+        );
+        let unique: std::collections::HashSet<_> = result.iter().collect();
+        assert_eq!(
+            unique.len(),
+            result.len(),
+            "each document should appear at most once"
+        );
+        assert_eq!(result.len(), 2);
+    }
+
+    #[test]
+    fn reciprocal_rank_fusion_respects_top_k() {
+        let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect();
+        let result = super::reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
+        assert_eq!(result.len(), 3, "result should be capped at top_k=3");
+    }
+
+    #[test]
+    fn reciprocal_rank_fusion_weights_affect_ranking() {
+        let doc_a = DocumentId::new(0, 0);
+        let doc_b = DocumentId::new(0, 1);
+        let result = super::reciprocal_rank_fusion(
+            vec![vec![doc_a, doc_b], vec![doc_b, doc_a]],
+            vec![10.0, 1.0],
+            2,
+        );
+        assert_eq!(
+            result[0], doc_a,
+            "higher-weight signal's top doc should rank first"
+        );
+    }
 }