1use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, HashSet, VecDeque};
8use std::sync::{Arc, RwLock};
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct PersonaNode {
13 pub persona_id: String,
15 pub entity_type: String,
17 pub relationships: HashMap<String, Vec<String>>,
21 #[serde(default)]
23 pub metadata: HashMap<String, serde_json::Value>,
24}
25
26impl PersonaNode {
27 pub fn new(persona_id: String, entity_type: String) -> Self {
29 Self {
30 persona_id,
31 entity_type,
32 relationships: HashMap::new(),
33 metadata: HashMap::new(),
34 }
35 }
36
37 pub fn add_relationship(&mut self, relationship_type: String, related_persona_id: String) {
39 self.relationships
40 .entry(relationship_type)
41 .or_default()
42 .push(related_persona_id);
43 }
44
45 pub fn get_related(&self, relationship_type: &str) -> Vec<String> {
47 self.relationships.get(relationship_type).cloned().unwrap_or_default()
48 }
49
50 pub fn get_relationship_types(&self) -> Vec<String> {
52 self.relationships.keys().cloned().collect()
53 }
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct Edge {
59 pub from: String,
61 pub to: String,
63 pub relationship_type: String,
65 #[serde(default = "default_edge_weight")]
67 pub weight: f64,
68}
69
70fn default_edge_weight() -> f64 {
71 1.0
72}
73
74#[derive(Debug, Clone)]
79pub struct PersonaGraph {
80 nodes: Arc<RwLock<HashMap<String, PersonaNode>>>,
82 edges: Arc<RwLock<HashMap<String, Vec<Edge>>>>,
84 reverse_edges: Arc<RwLock<HashMap<String, Vec<Edge>>>>,
86}
87
88impl PersonaGraph {
89 pub fn new() -> Self {
91 Self {
92 nodes: Arc::new(RwLock::new(HashMap::new())),
93 edges: Arc::new(RwLock::new(HashMap::new())),
94 reverse_edges: Arc::new(RwLock::new(HashMap::new())),
95 }
96 }
97
98 pub fn add_node(&self, node: PersonaNode) {
100 let mut nodes = self.nodes.write().expect("persona graph nodes write lock poisoned");
101 nodes.insert(node.persona_id.clone(), node);
102 }
103
104 pub fn get_node(&self, persona_id: &str) -> Option<PersonaNode> {
106 let nodes = self.nodes.read().expect("persona graph nodes read lock poisoned");
107 nodes.get(persona_id).cloned()
108 }
109
110 pub fn add_edge(&self, from: String, to: String, relationship_type: String) {
112 let to_clone = to.clone();
113 let edge = Edge {
114 from: from.clone(),
115 to: to_clone.clone(),
116 relationship_type: relationship_type.clone(),
117 weight: 1.0,
118 };
119
120 let mut edges = self.edges.write().expect("persona graph edges write lock poisoned");
122 edges.entry(from.clone()).or_default().push(edge.clone());
123
124 let mut reverse_edges = self
126 .reverse_edges
127 .write()
128 .expect("persona graph reverse_edges write lock poisoned");
129 reverse_edges.entry(to_clone.clone()).or_default().push(edge);
130
131 if let Some(node) = self.get_node(&from) {
133 let mut updated_node = node;
134 updated_node.add_relationship(relationship_type, to_clone);
135 self.add_node(updated_node);
136 }
137 }
138
139 pub fn get_edges_from(&self, persona_id: &str) -> Vec<Edge> {
141 let edges = self.edges.read().expect("persona graph edges read lock poisoned");
142 edges.get(persona_id).cloned().unwrap_or_default()
143 }
144
145 pub fn get_edges_to(&self, persona_id: &str) -> Vec<Edge> {
147 let reverse_edges = self
148 .reverse_edges
149 .read()
150 .expect("persona graph reverse_edges read lock poisoned");
151 reverse_edges.get(persona_id).cloned().unwrap_or_default()
152 }
153
154 pub fn find_related_bfs(
167 &self,
168 start_persona_id: &str,
169 relationship_types: Option<&[String]>,
170 max_depth: Option<usize>,
171 ) -> Vec<String> {
172 let mut visited = HashSet::new();
173 let mut queue = VecDeque::new();
174 let mut result = Vec::new();
175
176 queue.push_back((start_persona_id.to_string(), 0));
177 visited.insert(start_persona_id.to_string());
178
179 while let Some((current_id, depth)) = queue.pop_front() {
180 if let Some(max) = max_depth {
181 if depth >= max {
182 continue;
183 }
184 }
185
186 let edges = self.get_edges_from(¤t_id);
187 for edge in edges {
188 if let Some(types) = relationship_types {
190 if !types.contains(&edge.relationship_type) {
191 continue;
192 }
193 }
194
195 if !visited.contains(&edge.to) {
196 visited.insert(edge.to.clone());
197 result.push(edge.to.clone());
198 queue.push_back((edge.to.clone(), depth + 1));
199 }
200 }
201 }
202
203 result
204 }
205
206 pub fn find_related_dfs(
219 &self,
220 start_persona_id: &str,
221 relationship_types: Option<&[String]>,
222 max_depth: Option<usize>,
223 ) -> Vec<String> {
224 let mut visited = HashSet::new();
225 let mut result = Vec::new();
226
227 self.dfs_recursive(
228 start_persona_id,
229 relationship_types,
230 max_depth,
231 0,
232 &mut visited,
233 &mut result,
234 );
235
236 result
237 }
238
239 fn dfs_recursive(
241 &self,
242 current_id: &str,
243 relationship_types: Option<&[String]>,
244 max_depth: Option<usize>,
245 current_depth: usize,
246 visited: &mut HashSet<String>,
247 result: &mut Vec<String>,
248 ) {
249 if visited.contains(current_id) {
250 return;
251 }
252
253 if let Some(max) = max_depth {
254 if current_depth >= max {
255 return;
256 }
257 }
258
259 visited.insert(current_id.to_string());
260 if current_depth > 0 {
261 result.push(current_id.to_string());
263 }
264
265 let edges = self.get_edges_from(current_id);
266 for edge in edges {
267 if let Some(types) = relationship_types {
269 if !types.contains(&edge.relationship_type) {
270 continue;
271 }
272 }
273
274 self.dfs_recursive(
275 &edge.to,
276 relationship_types,
277 max_depth,
278 current_depth + 1,
279 visited,
280 result,
281 );
282 }
283 }
284
285 pub fn get_subgraph(&self, start_persona_id: &str) -> (Vec<PersonaNode>, Vec<Edge>) {
289 let related_ids = self.find_related_bfs(start_persona_id, None, None);
290 let mut all_ids = vec![start_persona_id.to_string()];
291 all_ids.extend(related_ids);
292
293 let nodes = self.nodes.read().expect("persona graph nodes read lock poisoned");
294 let edges = self.edges.read().expect("persona graph edges read lock poisoned");
295
296 let subgraph_nodes: Vec<PersonaNode> =
297 all_ids.iter().filter_map(|id| nodes.get(id).cloned()).collect();
298
299 let subgraph_edges: Vec<Edge> = all_ids
300 .iter()
301 .flat_map(|id| edges.get(id).cloned().unwrap_or_default())
302 .filter(|edge| all_ids.contains(&edge.to))
303 .collect();
304
305 (subgraph_nodes, subgraph_edges)
306 }
307
308 pub fn get_all_nodes(&self) -> Vec<PersonaNode> {
310 let nodes = self.nodes.read().expect("persona graph nodes read lock poisoned");
311 nodes.values().cloned().collect()
312 }
313
314 pub fn remove_node(&self, persona_id: &str) {
316 let mut nodes = self.nodes.write().expect("persona graph nodes write lock poisoned");
317 nodes.remove(persona_id);
318
319 let mut edges = self.edges.write().expect("persona graph edges write lock poisoned");
321 edges.remove(persona_id);
322
323 let mut reverse_edges = self
325 .reverse_edges
326 .write()
327 .expect("persona graph reverse_edges write lock poisoned");
328 reverse_edges.remove(persona_id);
329
330 for edges_list in edges.values_mut() {
332 edges_list.retain(|e| e.to != persona_id);
333 }
334 for edges_list in reverse_edges.values_mut() {
335 edges_list.retain(|e| e.from != persona_id);
336 }
337 }
338
339 pub fn clear(&self) {
341 let mut nodes = self.nodes.write().expect("persona graph nodes write lock poisoned");
342 nodes.clear();
343
344 let mut edges = self.edges.write().expect("persona graph edges write lock poisoned");
345 edges.clear();
346
347 let mut reverse_edges = self
348 .reverse_edges
349 .write()
350 .expect("persona graph reverse_edges write lock poisoned");
351 reverse_edges.clear();
352 }
353
354 pub fn get_stats(&self) -> GraphStats {
356 let nodes = self.nodes.read().expect("persona graph nodes read lock poisoned");
357 let edges = self.edges.read().expect("persona graph edges read lock poisoned");
358
359 let mut relationship_type_counts = HashMap::new();
360 for edges_list in edges.values() {
361 for edge in edges_list {
362 *relationship_type_counts.entry(edge.relationship_type.clone()).or_insert(0) += 1;
363 }
364 }
365
366 GraphStats {
367 node_count: nodes.len(),
368 edge_count: edges.values().map(|e| e.len()).sum(),
369 relationship_types: relationship_type_counts,
370 }
371 }
372
373 pub fn link_entity_types(
388 &self,
389 from_persona_id: &str,
390 from_entity_type: &str,
391 to_persona_id: &str,
392 to_entity_type: &str,
393 ) {
394 let relationship_type: String = match (from_entity_type, to_entity_type) {
396 ("user", "order") | ("user", "orders") => "has_orders".to_string(),
397 ("user", "account") | ("user", "accounts") => "has_accounts".to_string(),
398 ("user", "webhook") | ("user", "webhooks") => "has_webhooks".to_string(),
399 ("user", "tcp_message") | ("user", "tcp_messages") => "has_tcp_messages".to_string(),
400 ("order", "payment") | ("order", "payments") => "has_payments".to_string(),
401 ("account", "order") | ("account", "orders") => "has_orders".to_string(),
402 ("account", "payment") | ("account", "payments") => "has_payments".to_string(),
403 _ => {
404 format!("has_{}", to_entity_type.to_lowercase().trim_end_matches('s'))
406 }
407 };
408
409 if self.get_node(from_persona_id).is_none() {
411 let node = PersonaNode::new(from_persona_id.to_string(), from_entity_type.to_string());
412 self.add_node(node);
413 }
414
415 if self.get_node(to_persona_id).is_none() {
416 let node = PersonaNode::new(to_persona_id.to_string(), to_entity_type.to_string());
417 self.add_node(node);
418 }
419
420 self.add_edge(
422 from_persona_id.to_string(),
423 to_persona_id.to_string(),
424 relationship_type.to_string(),
425 );
426 }
427
428 pub fn find_related_by_entity_type(
441 &self,
442 start_persona_id: &str,
443 target_entity_type: &str,
444 relationship_type: Option<&str>,
445 ) -> Vec<String> {
446 let related_ids = if let Some(rel_type) = relationship_type {
447 let rel_types = vec![rel_type.to_string()];
448 self.find_related_bfs(start_persona_id, Some(&rel_types), Some(2))
449 } else {
450 self.find_related_bfs(start_persona_id, None, Some(2))
451 };
452
453 related_ids
455 .into_iter()
456 .filter_map(|persona_id| {
457 if let Some(node) = self.get_node(&persona_id) {
458 if node.entity_type.to_lowercase() == target_entity_type.to_lowercase() {
459 Some(persona_id)
460 } else {
461 None
462 }
463 } else {
464 None
465 }
466 })
467 .collect()
468 }
469
470 pub fn get_or_create_node_with_links(
481 &self,
482 persona_id: &str,
483 entity_type: &str,
484 related_entity_id: Option<&str>,
485 related_entity_type: Option<&str>,
486 ) -> PersonaNode {
487 let node = if let Some(existing) = self.get_node(persona_id) {
489 existing
490 } else {
491 let new_node = PersonaNode::new(persona_id.to_string(), entity_type.to_string());
492 self.add_node(new_node.clone());
493 new_node
494 };
495
496 if let (Some(related_id), Some(related_type)) = (related_entity_id, related_entity_type) {
498 self.link_entity_types(persona_id, entity_type, related_id, related_type);
499 }
500
501 node
502 }
503}
504
505impl Default for PersonaGraph {
506 fn default() -> Self {
507 Self::new()
508 }
509}
510
511#[derive(Debug, Clone, Serialize, Deserialize)]
513pub struct GraphStats {
514 pub node_count: usize,
516 pub edge_count: usize,
518 pub relationship_types: HashMap<String, usize>,
520}
521
522#[derive(Debug, Clone, Serialize, Deserialize)]
524pub struct GraphVisualization {
525 pub nodes: Vec<VisualizationNode>,
527 pub edges: Vec<VisualizationEdge>,
529}
530
531#[derive(Debug, Clone, Serialize, Deserialize)]
533pub struct VisualizationNode {
534 pub id: String,
536 pub entity_type: String,
538 pub label: String,
540 #[serde(skip_serializing_if = "Option::is_none")]
542 pub position: Option<(f64, f64)>,
543}
544
545#[derive(Debug, Clone, Serialize, Deserialize)]
547pub struct VisualizationEdge {
548 pub from: String,
550 pub to: String,
552 pub relationship_type: String,
554 pub label: String,
556}
557
558impl PersonaGraph {
559 pub fn to_visualization(&self) -> GraphVisualization {
561 let nodes = self.get_all_nodes();
562 let edges = self.edges.read().expect("persona graph edges read lock poisoned");
563
564 let vis_nodes: Vec<VisualizationNode> = nodes
565 .iter()
566 .map(|node| VisualizationNode {
567 id: node.persona_id.clone(),
568 entity_type: node.entity_type.clone(),
569 label: format!("{} ({})", node.persona_id, node.entity_type),
570 position: None,
571 })
572 .collect();
573
574 let vis_edges: Vec<VisualizationEdge> = edges
575 .values()
576 .flatten()
577 .map(|edge| VisualizationEdge {
578 from: edge.from.clone(),
579 to: edge.to.clone(),
580 relationship_type: edge.relationship_type.clone(),
581 label: edge.relationship_type.clone(),
582 })
583 .collect();
584
585 GraphVisualization {
586 nodes: vis_nodes,
587 edges: vis_edges,
588 }
589 }
590}
591
592#[cfg(test)]
593mod tests {
594 use super::*;
595
596 #[test]
601 fn test_persona_node_new() {
602 let node = PersonaNode::new("user-123".to_string(), "user".to_string());
603 assert_eq!(node.persona_id, "user-123");
604 assert_eq!(node.entity_type, "user");
605 assert!(node.relationships.is_empty());
606 assert!(node.metadata.is_empty());
607 }
608
609 #[test]
610 fn test_persona_node_add_relationship() {
611 let mut node = PersonaNode::new("user-123".to_string(), "user".to_string());
612 node.add_relationship("has_orders".to_string(), "order-1".to_string());
613 node.add_relationship("has_orders".to_string(), "order-2".to_string());
614
615 let related = node.get_related("has_orders");
616 assert_eq!(related.len(), 2);
617 assert!(related.contains(&"order-1".to_string()));
618 assert!(related.contains(&"order-2".to_string()));
619 }
620
621 #[test]
622 fn test_persona_node_get_related_empty() {
623 let node = PersonaNode::new("user-123".to_string(), "user".to_string());
624 let related = node.get_related("has_orders");
625 assert!(related.is_empty());
626 }
627
628 #[test]
629 fn test_persona_node_get_relationship_types() {
630 let mut node = PersonaNode::new("user-123".to_string(), "user".to_string());
631 node.add_relationship("has_orders".to_string(), "order-1".to_string());
632 node.add_relationship("has_payments".to_string(), "payment-1".to_string());
633
634 let types = node.get_relationship_types();
635 assert_eq!(types.len(), 2);
636 assert!(types.contains(&"has_orders".to_string()));
637 assert!(types.contains(&"has_payments".to_string()));
638 }
639
640 #[test]
641 fn test_persona_node_clone() {
642 let mut node = PersonaNode::new("user-123".to_string(), "user".to_string());
643 node.add_relationship("has_orders".to_string(), "order-1".to_string());
644
645 let cloned = node.clone();
646 assert_eq!(cloned.persona_id, node.persona_id);
647 assert_eq!(cloned.entity_type, node.entity_type);
648 assert_eq!(cloned.relationships, node.relationships);
649 }
650
651 #[test]
652 fn test_persona_node_debug() {
653 let node = PersonaNode::new("user-123".to_string(), "user".to_string());
654 let debug_str = format!("{:?}", node);
655 assert!(debug_str.contains("user-123"));
656 assert!(debug_str.contains("user"));
657 }
658
659 #[test]
660 fn test_persona_node_serialize_deserialize() {
661 let mut node = PersonaNode::new("user-123".to_string(), "user".to_string());
662 node.add_relationship("has_orders".to_string(), "order-1".to_string());
663
664 let json = serde_json::to_string(&node).unwrap();
665 let deserialized: PersonaNode = serde_json::from_str(&json).unwrap();
666
667 assert_eq!(deserialized.persona_id, "user-123");
668 assert_eq!(deserialized.entity_type, "user");
669 }
670
671 #[test]
676 fn test_edge_creation() {
677 let edge = Edge {
678 from: "user-123".to_string(),
679 to: "order-456".to_string(),
680 relationship_type: "has_orders".to_string(),
681 weight: 1.0,
682 };
683 assert_eq!(edge.from, "user-123");
684 assert_eq!(edge.to, "order-456");
685 assert_eq!(edge.relationship_type, "has_orders");
686 assert!((edge.weight - 1.0).abs() < f64::EPSILON);
687 }
688
689 #[test]
690 fn test_edge_clone() {
691 let edge = Edge {
692 from: "a".to_string(),
693 to: "b".to_string(),
694 relationship_type: "rel".to_string(),
695 weight: 2.5,
696 };
697 let cloned = edge.clone();
698 assert_eq!(cloned.from, edge.from);
699 assert!((cloned.weight - 2.5).abs() < f64::EPSILON);
700 }
701
702 #[test]
703 fn test_edge_debug() {
704 let edge = Edge {
705 from: "a".to_string(),
706 to: "b".to_string(),
707 relationship_type: "rel".to_string(),
708 weight: 1.0,
709 };
710 let debug_str = format!("{:?}", edge);
711 assert!(debug_str.contains("from"));
712 assert!(debug_str.contains("to"));
713 }
714
715 #[test]
716 fn test_edge_serialize_default_weight() {
717 let edge = Edge {
718 from: "a".to_string(),
719 to: "b".to_string(),
720 relationship_type: "rel".to_string(),
721 weight: 1.0,
722 };
723 let json = serde_json::to_string(&edge).unwrap();
724 let deserialized: Edge = serde_json::from_str(&json).unwrap();
725 assert!((deserialized.weight - 1.0).abs() < f64::EPSILON);
726 }
727
728 #[test]
733 fn test_persona_graph_new() {
734 let graph = PersonaGraph::new();
735 let stats = graph.get_stats();
736 assert_eq!(stats.node_count, 0);
737 assert_eq!(stats.edge_count, 0);
738 }
739
740 #[test]
741 fn test_persona_graph_default() {
742 let graph = PersonaGraph::default();
743 assert_eq!(graph.get_stats().node_count, 0);
744 }
745
746 #[test]
747 fn test_persona_graph_add_node() {
748 let graph = PersonaGraph::new();
749 let node = PersonaNode::new("user-123".to_string(), "user".to_string());
750 graph.add_node(node);
751
752 let stats = graph.get_stats();
753 assert_eq!(stats.node_count, 1);
754 }
755
756 #[test]
757 fn test_persona_graph_get_node() {
758 let graph = PersonaGraph::new();
759 let node = PersonaNode::new("user-123".to_string(), "user".to_string());
760 graph.add_node(node);
761
762 let retrieved = graph.get_node("user-123");
763 assert!(retrieved.is_some());
764 assert_eq!(retrieved.unwrap().persona_id, "user-123");
765 }
766
767 #[test]
768 fn test_persona_graph_get_node_not_found() {
769 let graph = PersonaGraph::new();
770 assert!(graph.get_node("nonexistent").is_none());
771 }
772
773 #[test]
774 fn test_persona_graph_add_edge() {
775 let graph = PersonaGraph::new();
776 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
777 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
778
779 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
780
781 let stats = graph.get_stats();
782 assert_eq!(stats.edge_count, 1);
783 }
784
785 #[test]
786 fn test_persona_graph_get_edges_from() {
787 let graph = PersonaGraph::new();
788 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
789 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
790 graph.add_node(PersonaNode::new("order-2".to_string(), "order".to_string()));
791
792 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
793 graph.add_edge("user-1".to_string(), "order-2".to_string(), "has_orders".to_string());
794
795 let edges = graph.get_edges_from("user-1");
796 assert_eq!(edges.len(), 2);
797 }
798
799 #[test]
800 fn test_persona_graph_get_edges_to() {
801 let graph = PersonaGraph::new();
802 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
803 graph.add_node(PersonaNode::new("user-2".to_string(), "user".to_string()));
804 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
805
806 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
807 graph.add_edge("user-2".to_string(), "order-1".to_string(), "has_orders".to_string());
808
809 let edges = graph.get_edges_to("order-1");
810 assert_eq!(edges.len(), 2);
811 }
812
813 #[test]
814 fn test_persona_graph_find_related_bfs() {
815 let graph = PersonaGraph::new();
816 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
817 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
818 graph.add_node(PersonaNode::new("payment-1".to_string(), "payment".to_string()));
819
820 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
821 graph.add_edge("order-1".to_string(), "payment-1".to_string(), "has_payments".to_string());
822
823 let related = graph.find_related_bfs("user-1", None, None);
824 assert_eq!(related.len(), 2);
825 assert!(related.contains(&"order-1".to_string()));
826 assert!(related.contains(&"payment-1".to_string()));
827 }
828
829 #[test]
830 fn test_persona_graph_find_related_bfs_with_depth_limit() {
831 let graph = PersonaGraph::new();
832 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
833 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
834 graph.add_node(PersonaNode::new("payment-1".to_string(), "payment".to_string()));
835
836 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
837 graph.add_edge("order-1".to_string(), "payment-1".to_string(), "has_payments".to_string());
838
839 let related = graph.find_related_bfs("user-1", None, Some(1));
840 assert_eq!(related.len(), 1);
841 assert!(related.contains(&"order-1".to_string()));
842 }
843
844 #[test]
845 fn test_persona_graph_find_related_bfs_with_type_filter() {
846 let graph = PersonaGraph::new();
847 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
848 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
849 graph.add_node(PersonaNode::new("account-1".to_string(), "account".to_string()));
850
851 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
852 graph.add_edge("user-1".to_string(), "account-1".to_string(), "has_accounts".to_string());
853
854 let filter = vec!["has_orders".to_string()];
855 let related = graph.find_related_bfs("user-1", Some(&filter), None);
856 assert_eq!(related.len(), 1);
857 assert!(related.contains(&"order-1".to_string()));
858 }
859
860 #[test]
861 fn test_persona_graph_find_related_dfs() {
862 let graph = PersonaGraph::new();
863 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
864 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
865 graph.add_node(PersonaNode::new("payment-1".to_string(), "payment".to_string()));
866
867 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
868 graph.add_edge("order-1".to_string(), "payment-1".to_string(), "has_payments".to_string());
869
870 let related = graph.find_related_dfs("user-1", None, None);
871 assert_eq!(related.len(), 2);
872 }
873
874 #[test]
875 fn test_persona_graph_find_related_dfs_with_depth_limit() {
876 let graph = PersonaGraph::new();
877 graph.add_node(PersonaNode::new("a".to_string(), "node".to_string()));
878 graph.add_node(PersonaNode::new("b".to_string(), "node".to_string()));
879 graph.add_node(PersonaNode::new("c".to_string(), "node".to_string()));
880 graph.add_node(PersonaNode::new("d".to_string(), "node".to_string()));
881
882 graph.add_edge("a".to_string(), "b".to_string(), "linked".to_string());
883 graph.add_edge("b".to_string(), "c".to_string(), "linked".to_string());
884 graph.add_edge("c".to_string(), "d".to_string(), "linked".to_string());
885
886 let related = graph.find_related_dfs("a", None, Some(2));
889 assert_eq!(related.len(), 1); assert!(related.contains(&"b".to_string()));
891 }
892
893 #[test]
894 fn test_persona_graph_get_subgraph() {
895 let graph = PersonaGraph::new();
896 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
897 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
898 graph.add_node(PersonaNode::new("isolated".to_string(), "node".to_string()));
899
900 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
901
902 let (nodes, edges) = graph.get_subgraph("user-1");
903 assert_eq!(nodes.len(), 2); assert_eq!(edges.len(), 1);
905 }
906
907 #[test]
908 fn test_persona_graph_get_all_nodes() {
909 let graph = PersonaGraph::new();
910 graph.add_node(PersonaNode::new("a".to_string(), "node".to_string()));
911 graph.add_node(PersonaNode::new("b".to_string(), "node".to_string()));
912 graph.add_node(PersonaNode::new("c".to_string(), "node".to_string()));
913
914 let nodes = graph.get_all_nodes();
915 assert_eq!(nodes.len(), 3);
916 }
917
918 #[test]
919 fn test_persona_graph_remove_node() {
920 let graph = PersonaGraph::new();
921 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
922 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
923
924 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
925
926 graph.remove_node("order-1");
927
928 assert!(graph.get_node("order-1").is_none());
929 assert_eq!(graph.get_edges_from("user-1").len(), 0);
930 }
931
932 #[test]
933 fn test_persona_graph_clear() {
934 let graph = PersonaGraph::new();
935 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
936 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
937 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
938
939 graph.clear();
940
941 let stats = graph.get_stats();
942 assert_eq!(stats.node_count, 0);
943 assert_eq!(stats.edge_count, 0);
944 }
945
946 #[test]
947 fn test_persona_graph_get_stats() {
948 let graph = PersonaGraph::new();
949 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
950 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
951 graph.add_node(PersonaNode::new("order-2".to_string(), "order".to_string()));
952
953 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
954 graph.add_edge("user-1".to_string(), "order-2".to_string(), "has_orders".to_string());
955
956 let stats = graph.get_stats();
957 assert_eq!(stats.node_count, 3);
958 assert_eq!(stats.edge_count, 2);
959 assert_eq!(*stats.relationship_types.get("has_orders").unwrap(), 2);
960 }
961
962 #[test]
967 fn test_persona_graph_link_entity_types_user_order() {
968 let graph = PersonaGraph::new();
969 graph.link_entity_types("user-1", "user", "order-1", "order");
970
971 let node = graph.get_node("user-1").unwrap();
972 assert_eq!(node.entity_type, "user");
973
974 let related = node.get_related("has_orders");
975 assert!(related.contains(&"order-1".to_string()));
976 }
977
978 #[test]
979 fn test_persona_graph_link_entity_types_order_payment() {
980 let graph = PersonaGraph::new();
981 graph.link_entity_types("order-1", "order", "payment-1", "payment");
982
983 let related = graph.get_node("order-1").unwrap().get_related("has_payments");
984 assert!(related.contains(&"payment-1".to_string()));
985 }
986
987 #[test]
988 fn test_persona_graph_link_entity_types_generic() {
989 let graph = PersonaGraph::new();
990 graph.link_entity_types("foo-1", "foo", "bar-1", "bars");
991
992 let node = graph.get_node("foo-1").unwrap();
993 let related = node.get_related("has_bar");
995 assert!(related.contains(&"bar-1".to_string()));
996 }
997
998 #[test]
999 fn test_persona_graph_link_entity_types_creates_nodes() {
1000 let graph = PersonaGraph::new();
1001 graph.link_entity_types("new-user", "user", "new-order", "order");
1002
1003 assert!(graph.get_node("new-user").is_some());
1004 assert!(graph.get_node("new-order").is_some());
1005 }
1006
1007 #[test]
1012 fn test_persona_graph_find_related_by_entity_type() {
1013 let graph = PersonaGraph::new();
1014 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
1015 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
1016 graph.add_node(PersonaNode::new("order-2".to_string(), "order".to_string()));
1017 graph.add_node(PersonaNode::new("payment-1".to_string(), "payment".to_string()));
1018
1019 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
1020 graph.add_edge("user-1".to_string(), "order-2".to_string(), "has_orders".to_string());
1021 graph.add_edge("user-1".to_string(), "payment-1".to_string(), "has_payments".to_string());
1022
1023 let orders = graph.find_related_by_entity_type("user-1", "order", None);
1024 assert_eq!(orders.len(), 2);
1025 }
1026
1027 #[test]
1028 fn test_persona_graph_find_related_by_entity_type_with_filter() {
1029 let graph = PersonaGraph::new();
1030 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
1031 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
1032
1033 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
1034
1035 let orders = graph.find_related_by_entity_type("user-1", "order", Some("has_orders"));
1036 assert_eq!(orders.len(), 1);
1037 }
1038
1039 #[test]
1044 fn test_persona_graph_get_or_create_node_new() {
1045 let graph = PersonaGraph::new();
1046 let node = graph.get_or_create_node_with_links("user-new", "user", None, None);
1047 assert_eq!(node.persona_id, "user-new");
1048 assert!(graph.get_node("user-new").is_some());
1049 }
1050
1051 #[test]
1052 fn test_persona_graph_get_or_create_node_existing() {
1053 let graph = PersonaGraph::new();
1054 let node1 = PersonaNode::new("user-existing".to_string(), "user".to_string());
1055 graph.add_node(node1);
1056
1057 let node2 = graph.get_or_create_node_with_links("user-existing", "user", None, None);
1058 assert_eq!(node2.persona_id, "user-existing");
1059 }
1060
1061 #[test]
1062 fn test_persona_graph_get_or_create_node_with_link() {
1063 let graph = PersonaGraph::new();
1064 let _node = graph.get_or_create_node_with_links(
1065 "user-link",
1066 "user",
1067 Some("order-link"),
1068 Some("order"),
1069 );
1070
1071 assert!(graph.get_node("user-link").is_some());
1072 assert!(graph.get_node("order-link").is_some());
1073 assert_eq!(graph.get_edges_from("user-link").len(), 1);
1074 }
1075
1076 #[test]
1081 fn test_graph_stats_clone() {
1082 let stats = GraphStats {
1083 node_count: 5,
1084 edge_count: 10,
1085 relationship_types: {
1086 let mut map = HashMap::new();
1087 map.insert("has_orders".to_string(), 5);
1088 map
1089 },
1090 };
1091 let cloned = stats.clone();
1092 assert_eq!(cloned.node_count, 5);
1093 assert_eq!(cloned.edge_count, 10);
1094 }
1095
1096 #[test]
1097 fn test_graph_stats_debug() {
1098 let stats = GraphStats {
1099 node_count: 3,
1100 edge_count: 2,
1101 relationship_types: HashMap::new(),
1102 };
1103 let debug_str = format!("{:?}", stats);
1104 assert!(debug_str.contains("node_count"));
1105 assert!(debug_str.contains("edge_count"));
1106 }
1107
1108 #[test]
1109 fn test_graph_stats_serialize() {
1110 let stats = GraphStats {
1111 node_count: 1,
1112 edge_count: 2,
1113 relationship_types: HashMap::new(),
1114 };
1115 let json = serde_json::to_string(&stats).unwrap();
1116 assert!(json.contains("node_count"));
1117 }
1118
1119 #[test]
1124 fn test_visualization_node_creation() {
1125 let node = VisualizationNode {
1126 id: "user-1".to_string(),
1127 entity_type: "user".to_string(),
1128 label: "User 1".to_string(),
1129 position: Some((0.0, 0.0)),
1130 };
1131 assert_eq!(node.id, "user-1");
1132 }
1133
1134 #[test]
1135 fn test_visualization_edge_creation() {
1136 let edge = VisualizationEdge {
1137 from: "a".to_string(),
1138 to: "b".to_string(),
1139 relationship_type: "linked".to_string(),
1140 label: "Linked".to_string(),
1141 };
1142 assert_eq!(edge.from, "a");
1143 }
1144
1145 #[test]
1146 fn test_persona_graph_to_visualization() {
1147 let graph = PersonaGraph::new();
1148 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
1149 graph.add_node(PersonaNode::new("order-1".to_string(), "order".to_string()));
1150 graph.add_edge("user-1".to_string(), "order-1".to_string(), "has_orders".to_string());
1151
1152 let viz = graph.to_visualization();
1153 assert_eq!(viz.nodes.len(), 2);
1154 assert_eq!(viz.edges.len(), 1);
1155 }
1156
1157 #[test]
1158 fn test_visualization_serialize() {
1159 let viz = GraphVisualization {
1160 nodes: vec![VisualizationNode {
1161 id: "test".to_string(),
1162 entity_type: "node".to_string(),
1163 label: "Test".to_string(),
1164 position: None,
1165 }],
1166 edges: vec![],
1167 };
1168 let json = serde_json::to_string(&viz).unwrap();
1169 assert!(json.contains("test"));
1170 }
1171
1172 #[test]
1177 fn test_persona_graph_handles_cycles() {
1178 let graph = PersonaGraph::new();
1179 graph.add_node(PersonaNode::new("a".to_string(), "node".to_string()));
1180 graph.add_node(PersonaNode::new("b".to_string(), "node".to_string()));
1181 graph.add_node(PersonaNode::new("c".to_string(), "node".to_string()));
1182
1183 graph.add_edge("a".to_string(), "b".to_string(), "linked".to_string());
1185 graph.add_edge("b".to_string(), "c".to_string(), "linked".to_string());
1186 graph.add_edge("c".to_string(), "a".to_string(), "linked".to_string());
1187
1188 let related = graph.find_related_bfs("a", None, None);
1190 assert_eq!(related.len(), 2); let related_dfs = graph.find_related_dfs("a", None, None);
1194 assert_eq!(related_dfs.len(), 2);
1195 }
1196
1197 #[test]
1202 fn test_persona_graph_clone() {
1203 let graph = PersonaGraph::new();
1204 graph.add_node(PersonaNode::new("user-1".to_string(), "user".to_string()));
1205
1206 let cloned = graph.clone();
1207 assert!(cloned.get_node("user-1").is_some());
1209 }
1210}