mockforge_world_state/
query.rs

1//! World State Query - Flexible querying of world state
2//!
3//! This module provides a query builder for filtering and searching
4//! world state snapshots.
5
6use crate::model::{NodeType, StateEdge, StateLayer, StateNode};
7use serde::{Deserialize, Serialize};
8use std::collections::HashSet;
9
10/// Query for filtering world state
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct WorldStateQuery {
13    /// Filter by node types
14    #[serde(default)]
15    pub node_types: Option<HashSet<NodeType>>,
16    /// Filter by layers
17    #[serde(default)]
18    pub layers: Option<HashSet<StateLayer>>,
19    /// Filter by node IDs
20    #[serde(default)]
21    pub node_ids: Option<HashSet<String>>,
22    /// Filter by relationship types
23    #[serde(default)]
24    pub relationship_types: Option<HashSet<String>>,
25    /// Include edges in results
26    #[serde(default = "default_true")]
27    pub include_edges: bool,
28    /// Maximum depth for traversal (for graph queries)
29    #[serde(default)]
30    pub max_depth: Option<usize>,
31}
32
33fn default_true() -> bool {
34    true
35}
36
37impl Default for WorldStateQuery {
38    fn default() -> Self {
39        Self {
40            node_types: None,
41            layers: None,
42            node_ids: None,
43            relationship_types: None,
44            include_edges: true,
45            max_depth: None,
46        }
47    }
48}
49
50impl WorldStateQuery {
51    /// Create a new empty query
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Filter by node types
57    pub fn with_node_types(mut self, types: HashSet<NodeType>) -> Self {
58        self.node_types = Some(types);
59        self
60    }
61
62    /// Filter by layers
63    pub fn with_layers(mut self, layers: HashSet<StateLayer>) -> Self {
64        self.layers = Some(layers);
65        self
66    }
67
68    /// Filter by node IDs
69    pub fn with_node_ids(mut self, ids: HashSet<String>) -> Self {
70        self.node_ids = Some(ids);
71        self
72    }
73
74    /// Filter by relationship types
75    pub fn with_relationship_types(mut self, types: HashSet<String>) -> Self {
76        self.relationship_types = Some(types);
77        self
78    }
79
80    /// Set whether to include edges
81    pub fn include_edges(mut self, include: bool) -> Self {
82        self.include_edges = include;
83        self
84    }
85
86    /// Set maximum traversal depth
87    pub fn with_max_depth(mut self, depth: usize) -> Self {
88        self.max_depth = Some(depth);
89        self
90    }
91
92    /// Check if a node matches this query
93    pub fn matches_node(&self, node: &StateNode) -> bool {
94        // Check node type filter
95        if let Some(ref types) = self.node_types {
96            if !types.contains(&node.node_type) {
97                return false;
98            }
99        }
100
101        // Check layer filter
102        if let Some(ref layers) = self.layers {
103            if !layers.contains(&node.layer) {
104                return false;
105            }
106        }
107
108        // Check node ID filter
109        if let Some(ref ids) = self.node_ids {
110            if !ids.contains(&node.id) {
111                return false;
112            }
113        }
114
115        true
116    }
117
118    /// Check if an edge matches this query
119    pub fn matches_edge(&self, edge: &StateEdge) -> bool {
120        // Check relationship type filter
121        if let Some(ref types) = self.relationship_types {
122            if !types.contains(&edge.relationship_type) {
123                return false;
124            }
125        }
126
127        true
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use chrono::Utc;
135
136    fn create_test_node(id: &str, node_type: NodeType, layer: StateLayer) -> StateNode {
137        StateNode {
138            id: id.to_string(),
139            label: format!("Test {}", id),
140            node_type,
141            layer,
142            state: None,
143            properties: std::collections::HashMap::new(),
144            created_at: Utc::now(),
145            updated_at: Utc::now(),
146        }
147    }
148
149    #[test]
150    fn test_query_new() {
151        let query = WorldStateQuery::new();
152        assert!(query.node_types.is_none());
153        assert!(query.layers.is_none());
154        assert!(query.node_ids.is_none());
155        assert!(query.relationship_types.is_none());
156        assert!(query.include_edges);
157        assert!(query.max_depth.is_none());
158    }
159
160    #[test]
161    fn test_query_default() {
162        let query = WorldStateQuery::default();
163        assert!(query.include_edges);
164    }
165
166    #[test]
167    fn test_query_with_node_types() {
168        let mut types = HashSet::new();
169        types.insert(NodeType::Persona);
170        types.insert(NodeType::Entity);
171
172        let query = WorldStateQuery::new().with_node_types(types.clone());
173        assert_eq!(query.node_types, Some(types));
174    }
175
176    #[test]
177    fn test_query_with_layers() {
178        let mut layers = HashSet::new();
179        layers.insert(StateLayer::Personas);
180        layers.insert(StateLayer::Lifecycle);
181
182        let query = WorldStateQuery::new().with_layers(layers.clone());
183        assert_eq!(query.layers, Some(layers));
184    }
185
186    #[test]
187    fn test_query_with_node_ids() {
188        let mut ids = HashSet::new();
189        ids.insert("node1".to_string());
190        ids.insert("node2".to_string());
191
192        let query = WorldStateQuery::new().with_node_ids(ids.clone());
193        assert_eq!(query.node_ids, Some(ids));
194    }
195
196    #[test]
197    fn test_query_with_relationship_types() {
198        let mut types = HashSet::new();
199        types.insert("owns".to_string());
200        types.insert("references".to_string());
201
202        let query = WorldStateQuery::new().with_relationship_types(types.clone());
203        assert_eq!(query.relationship_types, Some(types));
204    }
205
206    #[test]
207    fn test_query_include_edges() {
208        let query = WorldStateQuery::new().include_edges(false);
209        assert!(!query.include_edges);
210    }
211
212    #[test]
213    fn test_query_with_max_depth() {
214        let query = WorldStateQuery::new().with_max_depth(3);
215        assert_eq!(query.max_depth, Some(3));
216    }
217
218    #[test]
219    fn test_query_builder_chaining() {
220        let mut types = HashSet::new();
221        types.insert(NodeType::Persona);
222
223        let mut layers = HashSet::new();
224        layers.insert(StateLayer::Personas);
225
226        let query = WorldStateQuery::new()
227            .with_node_types(types)
228            .with_layers(layers)
229            .include_edges(false)
230            .with_max_depth(5);
231
232        assert!(query.node_types.is_some());
233        assert!(query.layers.is_some());
234        assert!(!query.include_edges);
235        assert_eq!(query.max_depth, Some(5));
236    }
237
238    #[test]
239    fn test_matches_node_no_filters() {
240        let query = WorldStateQuery::new();
241        let node = create_test_node("node1", NodeType::Persona, StateLayer::Personas);
242        assert!(query.matches_node(&node));
243    }
244
245    #[test]
246    fn test_matches_node_by_type() {
247        let mut types = HashSet::new();
248        types.insert(NodeType::Persona);
249
250        let query = WorldStateQuery::new().with_node_types(types);
251
252        let persona_node = create_test_node("node1", NodeType::Persona, StateLayer::Personas);
253        let entity_node = create_test_node("node2", NodeType::Entity, StateLayer::Lifecycle);
254
255        assert!(query.matches_node(&persona_node));
256        assert!(!query.matches_node(&entity_node));
257    }
258
259    #[test]
260    fn test_matches_node_by_layer() {
261        let mut layers = HashSet::new();
262        layers.insert(StateLayer::Personas);
263
264        let query = WorldStateQuery::new().with_layers(layers);
265
266        let persona_node = create_test_node("node1", NodeType::Persona, StateLayer::Personas);
267        let lifecycle_node = create_test_node("node2", NodeType::Entity, StateLayer::Lifecycle);
268
269        assert!(query.matches_node(&persona_node));
270        assert!(!query.matches_node(&lifecycle_node));
271    }
272
273    #[test]
274    fn test_matches_node_by_id() {
275        let mut ids = HashSet::new();
276        ids.insert("node1".to_string());
277        ids.insert("node2".to_string());
278
279        let query = WorldStateQuery::new().with_node_ids(ids);
280
281        let matching_node = create_test_node("node1", NodeType::Persona, StateLayer::Personas);
282        let non_matching_node = create_test_node("node3", NodeType::Persona, StateLayer::Personas);
283
284        assert!(query.matches_node(&matching_node));
285        assert!(!query.matches_node(&non_matching_node));
286    }
287
288    #[test]
289    fn test_matches_node_multiple_filters() {
290        let mut types = HashSet::new();
291        types.insert(NodeType::Persona);
292
293        let mut layers = HashSet::new();
294        layers.insert(StateLayer::Personas);
295
296        let query = WorldStateQuery::new().with_node_types(types).with_layers(layers);
297
298        // Matches both filters
299        let matching = create_test_node("node1", NodeType::Persona, StateLayer::Personas);
300        assert!(query.matches_node(&matching));
301
302        // Matches type but not layer
303        let wrong_layer = create_test_node("node2", NodeType::Persona, StateLayer::Lifecycle);
304        assert!(!query.matches_node(&wrong_layer));
305
306        // Matches layer but not type
307        let wrong_type = create_test_node("node3", NodeType::Entity, StateLayer::Personas);
308        assert!(!query.matches_node(&wrong_type));
309    }
310
311    #[test]
312    fn test_matches_edge_no_filters() {
313        let query = WorldStateQuery::new();
314        let edge = StateEdge::new("a".to_string(), "b".to_string(), "owns".to_string());
315        assert!(query.matches_edge(&edge));
316    }
317
318    #[test]
319    fn test_matches_edge_by_relationship_type() {
320        let mut types = HashSet::new();
321        types.insert("owns".to_string());
322
323        let query = WorldStateQuery::new().with_relationship_types(types);
324
325        let matching_edge = StateEdge::new("a".to_string(), "b".to_string(), "owns".to_string());
326        let non_matching_edge =
327            StateEdge::new("a".to_string(), "b".to_string(), "references".to_string());
328
329        assert!(query.matches_edge(&matching_edge));
330        assert!(!query.matches_edge(&non_matching_edge));
331    }
332
333    #[test]
334    fn test_query_serialize() {
335        let mut types = HashSet::new();
336        types.insert(NodeType::Persona);
337
338        let query = WorldStateQuery::new().with_node_types(types);
339        let json = serde_json::to_string(&query).unwrap();
340        assert!(json.contains("\"persona\""));
341    }
342
343    #[test]
344    fn test_query_deserialize() {
345        let json = r#"{
346            "node_types": ["persona"],
347            "layers": ["personas"],
348            "include_edges": false,
349            "max_depth": 3
350        }"#;
351        let query: WorldStateQuery = serde_json::from_str(json).unwrap();
352        assert!(query.node_types.is_some());
353        assert!(query.layers.is_some());
354        assert!(!query.include_edges);
355        assert_eq!(query.max_depth, Some(3));
356    }
357
358    #[test]
359    fn test_query_clone() {
360        let mut types = HashSet::new();
361        types.insert(NodeType::Persona);
362
363        let query = WorldStateQuery::new().with_node_types(types).with_max_depth(5);
364
365        let cloned = query.clone();
366        assert_eq!(query.node_types, cloned.node_types);
367        assert_eq!(query.max_depth, cloned.max_depth);
368    }
369}