Skip to main content

mockforge_data/
persona_graph.rs

1//! Persona Graph & Relationship Management
2//!
3//! This module provides graph-based relationship management for personas,
4//! enabling coherent persona switching across related entities (user → orders → payments → support tickets).
5
6use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, HashSet, VecDeque};
8use std::sync::{Arc, RwLock};
9
10/// Represents a node in the persona graph
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct PersonaNode {
13    /// Persona ID
14    pub persona_id: String,
15    /// Entity type (e.g., "user", "order", "payment", "support_ticket")
16    pub entity_type: String,
17    /// Relationships from this persona to others
18    /// Key: relationship type (e.g., "has_orders", "has_payments")
19    /// Value: List of related persona IDs
20    pub relationships: HashMap<String, Vec<String>>,
21    /// Additional metadata for the node
22    #[serde(default)]
23    pub metadata: HashMap<String, serde_json::Value>,
24}
25
26impl PersonaNode {
27    /// Create a new persona node
28    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    /// Add a relationship to another persona
38    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    /// Get all related personas for a relationship type
46    pub fn get_related(&self, relationship_type: &str) -> Vec<String> {
47        self.relationships.get(relationship_type).cloned().unwrap_or_default()
48    }
49
50    /// Get all relationship types for this node
51    pub fn get_relationship_types(&self) -> Vec<String> {
52        self.relationships.keys().cloned().collect()
53    }
54}
55
56/// Edge in the persona graph
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct Edge {
59    /// Source persona ID
60    pub from: String,
61    /// Target persona ID
62    pub to: String,
63    /// Relationship type
64    pub relationship_type: String,
65    /// Edge weight (for weighted traversals, default 1.0)
66    #[serde(default = "default_edge_weight")]
67    pub weight: f64,
68}
69
70fn default_edge_weight() -> f64 {
71    1.0
72}
73
74/// Persona graph for managing entity relationships
75///
76/// Maintains a graph structure of personas and their relationships,
77/// enabling coherent persona switching across related entities.
78#[derive(Debug, Clone)]
79pub struct PersonaGraph {
80    /// Graph nodes indexed by persona ID
81    nodes: Arc<RwLock<HashMap<String, PersonaNode>>>,
82    /// Graph edges indexed by source persona ID
83    edges: Arc<RwLock<HashMap<String, Vec<Edge>>>>,
84    /// Reverse edges for efficient backward traversal
85    reverse_edges: Arc<RwLock<HashMap<String, Vec<Edge>>>>,
86}
87
88impl PersonaGraph {
89    /// Create a new empty persona graph
90    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    /// Add a persona node to the graph
99    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    /// Get a node by persona ID
105    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    /// Add an edge between two personas
111    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        // Add forward edge
121        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        // Add reverse edge
125        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        // Update node relationships
132        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    /// Get all edges from a persona
140    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    /// Get all edges to a persona
146    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    /// Find all related personas using BFS traversal
155    ///
156    /// Traverses the graph starting from the given persona ID,
157    /// following relationships of the specified types.
158    ///
159    /// # Arguments
160    /// * `start_persona_id` - Starting persona ID
161    /// * `relationship_types` - Optional filter for relationship types to follow
162    /// * `max_depth` - Maximum traversal depth (None = unlimited)
163    ///
164    /// # Returns
165    /// Vector of persona IDs reachable from the start persona
166    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(&current_id);
187            for edge in edges {
188                // Filter by relationship type if specified
189                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    /// Find all related personas using DFS traversal
207    ///
208    /// Traverses the graph starting from the given persona ID,
209    /// following relationships of the specified types.
210    ///
211    /// # Arguments
212    /// * `start_persona_id` - Starting persona ID
213    /// * `relationship_types` - Optional filter for relationship types to follow
214    /// * `max_depth` - Maximum traversal depth (None = unlimited)
215    ///
216    /// # Returns
217    /// Vector of persona IDs reachable from the start persona
218    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    /// Recursive helper for DFS traversal
240    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            // Don't include the start node in results
262            result.push(current_id.to_string());
263        }
264
265        let edges = self.get_edges_from(current_id);
266        for edge in edges {
267            // Filter by relationship type if specified
268            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    /// Get the entire subgraph starting from a persona
286    ///
287    /// Returns all nodes and edges reachable from the start persona.
288    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    /// Get all nodes in the graph
309    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    /// Remove a node and all its edges
315    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        // Remove forward edges
320        let mut edges = self.edges.write().expect("persona graph edges write lock poisoned");
321        edges.remove(persona_id);
322
323        // Remove reverse edges
324        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        // Remove edges pointing to this node
331        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    /// Clear the entire graph
340    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    /// Get graph statistics
355    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    /// Link personas across entity types automatically
374    ///
375    /// Creates relationships between personas based on common entity type patterns:
376    /// - user → has_orders → order
377    /// - user → has_accounts → account
378    /// - order → has_payments → payment
379    /// - user → has_webhooks → webhook
380    /// - user → has_tcp_messages → tcp_message
381    ///
382    /// # Arguments
383    /// * `from_persona_id` - Source persona ID
384    /// * `from_entity_type` - Source entity type (e.g., "user", "order")
385    /// * `to_persona_id` - Target persona ID
386    /// * `to_entity_type` - Target entity type (e.g., "order", "payment")
387    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        // Determine relationship type based on entity types
395        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                // Generic relationship: from_entity_type -> to_entity_type
405                format!("has_{}", to_entity_type.to_lowercase().trim_end_matches('s'))
406            }
407        };
408
409        // Ensure both nodes exist
410        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        // Add the edge
421        self.add_edge(
422            from_persona_id.to_string(),
423            to_persona_id.to_string(),
424            relationship_type.to_string(),
425        );
426    }
427
428    /// Find all related personas of a specific entity type
429    ///
430    /// Traverses the graph to find all personas of the specified entity type
431    /// that are related to the starting persona.
432    ///
433    /// # Arguments
434    /// * `start_persona_id` - Starting persona ID
435    /// * `target_entity_type` - Entity type to find (e.g., "order", "payment")
436    /// * `relationship_type` - Optional relationship type filter (e.g., "has_orders")
437    ///
438    /// # Returns
439    /// Vector of persona IDs matching the criteria
440    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        // Filter by entity type
454        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    /// Get or create a persona node and link it to related entities
471    ///
472    /// This is a convenience method that creates a persona node if it doesn't exist
473    /// and automatically establishes relationships based on entity type patterns.
474    ///
475    /// # Arguments
476    /// * `persona_id` - Persona ID
477    /// * `entity_type` - Entity type (e.g., "user", "order", "payment")
478    /// * `related_entity_id` - Optional related entity ID to link to
479    /// * `related_entity_type` - Optional related entity type
480    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        // Get or create the node
488        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        // Link to related entity if provided
497        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/// Graph statistics
512#[derive(Debug, Clone, Serialize, Deserialize)]
513pub struct GraphStats {
514    /// Number of nodes in the graph
515    pub node_count: usize,
516    /// Number of edges in the graph
517    pub edge_count: usize,
518    /// Count of edges by relationship type
519    pub relationship_types: HashMap<String, usize>,
520}
521
522/// Graph visualization data structure
523#[derive(Debug, Clone, Serialize, Deserialize)]
524pub struct GraphVisualization {
525    /// Nodes in the graph
526    pub nodes: Vec<VisualizationNode>,
527    /// Edges in the graph
528    pub edges: Vec<VisualizationEdge>,
529}
530
531/// Node for visualization
532#[derive(Debug, Clone, Serialize, Deserialize)]
533pub struct VisualizationNode {
534    /// Persona ID
535    pub id: String,
536    /// Entity type
537    pub entity_type: String,
538    /// Display label
539    pub label: String,
540    /// Node position (for layout algorithms)
541    #[serde(skip_serializing_if = "Option::is_none")]
542    pub position: Option<(f64, f64)>,
543}
544
545/// Edge for visualization
546#[derive(Debug, Clone, Serialize, Deserialize)]
547pub struct VisualizationEdge {
548    /// Source persona ID
549    pub from: String,
550    /// Target persona ID
551    pub to: String,
552    /// Relationship type
553    pub relationship_type: String,
554    /// Display label
555    pub label: String,
556}
557
558impl PersonaGraph {
559    /// Generate visualization data for the graph
560    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    // =========================================================================
597    // PersonaNode tests
598    // =========================================================================
599
600    #[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    // =========================================================================
672    // Edge tests
673    // =========================================================================
674
675    #[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    // =========================================================================
729    // PersonaGraph tests
730    // =========================================================================
731
732    #[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        // DFS implementation: max_depth=2 means we can go 0->1->2, but depth 2 is the cutoff
887        // So we get nodes at depth 1 only (b), not c at depth 2
888        let related = graph.find_related_dfs("a", None, Some(2));
889        assert_eq!(related.len(), 1); // only b, depth 2 is the cutoff
890        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); // user-1 and order-1
904        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    // =========================================================================
963    // Link entity types tests
964    // =========================================================================
965
966    #[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        // Generic relationship: has_bar (from "bars" -> "bar")
994        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    // =========================================================================
1008    // Find related by entity type tests
1009    // =========================================================================
1010
1011    #[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    // =========================================================================
1040    // Get or create node with links tests
1041    // =========================================================================
1042
1043    #[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    // =========================================================================
1077    // GraphStats tests
1078    // =========================================================================
1079
1080    #[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    // =========================================================================
1120    // Visualization tests
1121    // =========================================================================
1122
1123    #[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    // =========================================================================
1173    // Cycle detection tests
1174    // =========================================================================
1175
1176    #[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        // Create a cycle: a -> b -> c -> a
1184        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        // BFS should not loop infinitely
1189        let related = graph.find_related_bfs("a", None, None);
1190        assert_eq!(related.len(), 2); // b and c (not a again)
1191
1192        // DFS should not loop infinitely
1193        let related_dfs = graph.find_related_dfs("a", None, None);
1194        assert_eq!(related_dfs.len(), 2);
1195    }
1196
1197    // =========================================================================
1198    // Clone tests
1199    // =========================================================================
1200
1201    #[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        // Both graphs share the same underlying data via Arc
1208        assert!(cloned.get_node("user-1").is_some());
1209    }
1210}