arrow_graph/graph/
tests.rs

1#[cfg(test)]
2mod tests {
3    use crate::ArrowGraph;
4    use arrow::record_batch::RecordBatch;
5    use arrow::array::{StringArray, Float64Array};
6    use arrow::datatypes::{DataType, Field, Schema};
7    use std::sync::Arc;
8
9    fn create_test_nodes() -> RecordBatch {
10        let schema = Arc::new(Schema::new(vec![
11            Field::new("id", DataType::Utf8, false),
12            Field::new("label", DataType::Utf8, true),
13        ]));
14
15        let id_array = StringArray::from(vec!["A", "B", "C", "D"]);
16        let label_array = StringArray::from(vec![
17            Some("Node A"), 
18            Some("Node B"), 
19            Some("Node C"), 
20            Some("Node D")
21        ]);
22
23        RecordBatch::try_new(
24            schema,
25            vec![
26                Arc::new(id_array) as Arc<dyn arrow::array::Array>, 
27                Arc::new(label_array) as Arc<dyn arrow::array::Array>
28            ],
29        ).unwrap()
30    }
31
32    fn create_test_edges() -> RecordBatch {
33        let schema = Arc::new(Schema::new(vec![
34            Field::new("source", DataType::Utf8, false),
35            Field::new("target", DataType::Utf8, false),
36            Field::new("weight", DataType::Float64, true),
37        ]));
38
39        let source_array = StringArray::from(vec!["A", "A", "B", "C"]);
40        let target_array = StringArray::from(vec!["B", "C", "C", "D"]);
41        let weight_array = Float64Array::from(vec![
42            Some(1.0), 
43            Some(2.0), 
44            Some(1.5), 
45            Some(0.5)
46        ]);
47
48        RecordBatch::try_new(
49            schema,
50            vec![
51                Arc::new(source_array) as Arc<dyn arrow::array::Array>, 
52                Arc::new(target_array) as Arc<dyn arrow::array::Array>, 
53                Arc::new(weight_array) as Arc<dyn arrow::array::Array>
54            ],
55        ).unwrap()
56    }
57
58    #[test]
59    fn test_graph_creation() {
60        let nodes = create_test_nodes();
61        let edges = create_test_edges();
62        
63        let graph = ArrowGraph::new(nodes, edges).unwrap();
64        
65        assert_eq!(graph.node_count(), 4);
66        assert_eq!(graph.edge_count(), 4);
67        
68        // Test adjacency
69        assert!(graph.has_node("A"));
70        assert!(graph.has_node("B"));
71        assert!(graph.has_node("C"));
72        assert!(graph.has_node("D"));
73        assert!(!graph.has_node("E"));
74        
75        // Test neighbors
76        let a_neighbors = graph.neighbors("A").unwrap();
77        assert_eq!(a_neighbors.len(), 2);
78        assert!(a_neighbors.contains(&"B".to_string()));
79        assert!(a_neighbors.contains(&"C".to_string()));
80        
81        // Test edge weights
82        assert_eq!(graph.edge_weight("A", "B"), Some(1.0));
83        assert_eq!(graph.edge_weight("A", "C"), Some(2.0));
84        assert_eq!(graph.edge_weight("A", "D"), None);
85    }
86
87    #[test]
88    fn test_graph_from_edges_only() {
89        let edges = create_test_edges();
90        let graph = ArrowGraph::from_edges(edges).unwrap();
91        
92        // Should infer 4 nodes from edges
93        assert_eq!(graph.node_count(), 4);
94        assert_eq!(graph.edge_count(), 4);
95        
96        assert!(graph.has_node("A"));
97        assert!(graph.has_node("B"));
98        assert!(graph.has_node("C"));
99        assert!(graph.has_node("D"));
100    }
101
102    #[test]
103    fn test_graph_density() {
104        let nodes = create_test_nodes();
105        let edges = create_test_edges();
106        let graph = ArrowGraph::new(nodes, edges).unwrap();
107        
108        // 4 nodes, 4 edges, max possible = 4 * 3 = 12
109        // density = 4/12 = 1/3 ≈ 0.333
110        let density = graph.density();
111        assert!((density - 0.3333333333333333).abs() < 0.0001);
112    }
113
114    #[test]
115    fn test_predecessors() {
116        let nodes = create_test_nodes();
117        let edges = create_test_edges();
118        let graph = ArrowGraph::new(nodes, edges).unwrap();
119        
120        // C has predecessors A and B
121        let c_predecessors = graph.predecessors("C").unwrap();
122        assert_eq!(c_predecessors.len(), 2);
123        assert!(c_predecessors.contains(&"A".to_string()));
124        assert!(c_predecessors.contains(&"B".to_string()));
125        
126        // A has no predecessors
127        let a_predecessors = graph.predecessors("A").unwrap();
128        assert_eq!(a_predecessors.len(), 0);
129    }
130}
arrow_graph/graph/tests.rs

arrow_graph/graph/
tests.rs
