1use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6
7use super::nodes::{TestNode, TestNodeId};
8use crate::storage::{keys, StorageBackend, StorageError, StorageResult, SCHEMA_VERSION};
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct InventoryMeta {
13 pub hash: String,
15
16 pub collected_at: u64,
18
19 pub node_count: usize,
21
22 pub schema_version: u32,
24}
25
26impl Default for InventoryMeta {
27 fn default() -> Self {
28 Self {
29 hash: String::new(),
30 collected_at: 0,
31 node_count: 0,
32 schema_version: SCHEMA_VERSION,
33 }
34 }
35}
36
37#[derive(Debug, Clone, Default)]
39pub struct Inventory {
40 nodes: HashMap<TestNodeId, TestNode>,
42
43 meta: InventoryMeta,
45}
46
47impl Inventory {
48 pub fn new() -> Self {
50 Self::default()
51 }
52
53 pub fn add(&mut self, node: TestNode) {
55 self.nodes.insert(node.node_id.clone(), node);
56 self.meta.node_count = self.nodes.len();
57 }
58
59 pub fn get(&self, node_id: &str) -> Option<&TestNode> {
61 self.nodes.get(node_id)
62 }
63
64 pub fn get_mut(&mut self, node_id: &str) -> Option<&mut TestNode> {
66 self.nodes.get_mut(node_id)
67 }
68
69 pub fn remove(&mut self, node_id: &str) -> Option<TestNode> {
71 let result = self.nodes.remove(node_id);
72 self.meta.node_count = self.nodes.len();
73 result
74 }
75
76 pub fn contains(&self, node_id: &str) -> bool {
78 self.nodes.contains_key(node_id)
79 }
80
81 pub fn len(&self) -> usize {
83 self.nodes.len()
84 }
85
86 pub fn is_empty(&self) -> bool {
88 self.nodes.is_empty()
89 }
90
91 pub fn iter(&self) -> impl Iterator<Item = &TestNode> {
93 self.nodes.values()
94 }
95
96 pub fn node_ids(&self) -> Vec<&TestNodeId> {
98 self.nodes.keys().collect()
99 }
100
101 pub fn meta(&self) -> &InventoryMeta {
103 &self.meta
104 }
105
106 pub fn update_meta(&mut self, hash: String, collected_at: u64) {
108 self.meta.hash = hash;
109 self.meta.collected_at = collected_at;
110 self.meta.node_count = self.nodes.len();
111 self.meta.schema_version = SCHEMA_VERSION;
112 }
113
114 pub fn filter_by_keyword(&self, expr: &str) -> Vec<&TestNode> {
116 self.nodes
117 .values()
118 .filter(|node| node.matches_keyword(expr))
119 .collect()
120 }
121
122 pub fn filter_by_marker(&self, expr: &str) -> Vec<&TestNode> {
124 self.nodes
125 .values()
126 .filter(|node| node.matches_marker(expr))
127 .collect()
128 }
129
130 pub fn filter(&self, keyword: Option<&str>, marker: Option<&str>) -> Vec<&TestNode> {
132 self.nodes
133 .values()
134 .filter(|node| {
135 keyword.map_or(true, |k| node.matches_keyword(k))
136 && marker.map_or(true, |m| node.matches_marker(m))
137 })
138 .collect()
139 }
140
141 pub fn sorted_by_duration(&self) -> Vec<&TestNode> {
143 let mut nodes: Vec<_> = self.nodes.values().collect();
144 nodes.sort_by(|a, b| {
145 b.avg_duration_ms
146 .unwrap_or(0)
147 .cmp(&a.avg_duration_ms.unwrap_or(0))
148 });
149 nodes
150 }
151
152 pub fn sorted_by_failure_rate(&self) -> Vec<&TestNode> {
154 let mut nodes: Vec<_> = self.nodes.values().collect();
155 nodes.sort_by(|a, b| {
156 b.failure_rate()
157 .partial_cmp(&a.failure_rate())
158 .unwrap_or(std::cmp::Ordering::Equal)
159 });
160 nodes
161 }
162
163 pub fn save<S: StorageBackend>(&self, storage: &S, context_id: &str) -> StorageResult<()> {
165 for (node_id, node) in &self.nodes {
167 let key = format!(
168 "{}{}:{}",
169 String::from_utf8_lossy(keys::INVENTORY),
170 context_id,
171 node_id
172 );
173 let value =
174 rmp_serde::to_vec(node).map_err(|e| StorageError::Serialization(e.to_string()))?;
175 storage.set(key.as_bytes(), &value)?;
176 }
177
178 let meta_key = format!(
180 "{}{}:_meta",
181 String::from_utf8_lossy(keys::CONTEXT),
182 context_id
183 );
184 let meta_value = rmp_serde::to_vec(&self.meta)
185 .map_err(|e| StorageError::Serialization(e.to_string()))?;
186 storage.set(meta_key.as_bytes(), &meta_value)?;
187
188 let schema_value = SCHEMA_VERSION.to_le_bytes();
190 storage.set(keys::SCHEMA, &schema_value)?;
191
192 storage.flush()?;
193 Ok(())
194 }
195
196 pub fn load<S: StorageBackend>(storage: &S, context_id: &str) -> StorageResult<Self> {
198 if let Some(version_bytes) = storage.get(keys::SCHEMA)? {
200 if version_bytes.len() >= 4 {
201 let version = u32::from_le_bytes([
202 version_bytes[0],
203 version_bytes[1],
204 version_bytes[2],
205 version_bytes[3],
206 ]);
207 if version != SCHEMA_VERSION {
208 return Err(StorageError::Corrupted(format!(
209 "Schema version mismatch: expected {}, found {}",
210 SCHEMA_VERSION, version
211 )));
212 }
213 }
214 }
215
216 let mut inventory = Self::new();
217
218 let meta_key = format!(
220 "{}{}:_meta",
221 String::from_utf8_lossy(keys::CONTEXT),
222 context_id
223 );
224 if let Some(meta_bytes) = storage.get(meta_key.as_bytes())? {
225 inventory.meta = rmp_serde::from_slice(&meta_bytes)
226 .map_err(|e| StorageError::Serialization(e.to_string()))?;
227 }
228
229 let prefix = format!(
231 "{}{}:",
232 String::from_utf8_lossy(keys::INVENTORY),
233 context_id
234 );
235 let entries = storage.scan_prefix(prefix.as_bytes())?;
236
237 for (_key, value) in entries {
238 let node: TestNode = rmp_serde::from_slice(&value)
239 .map_err(|e| StorageError::Serialization(e.to_string()))?;
240 inventory.nodes.insert(node.node_id.clone(), node);
241 }
242
243 inventory.meta.node_count = inventory.nodes.len();
244 Ok(inventory)
245 }
246
247 pub fn clear_storage<S: StorageBackend>(storage: &S, context_id: &str) -> StorageResult<()> {
249 let prefix = format!(
251 "{}{}:",
252 String::from_utf8_lossy(keys::INVENTORY),
253 context_id
254 );
255 let entries = storage.scan_prefix(prefix.as_bytes())?;
256 for (key, _) in entries {
257 storage.delete(&key)?;
258 }
259
260 let meta_key = format!(
262 "{}{}:_meta",
263 String::from_utf8_lossy(keys::CONTEXT),
264 context_id
265 );
266 storage.delete(meta_key.as_bytes())?;
267
268 storage.flush()?;
269 Ok(())
270 }
271}
272
273#[cfg(test)]
274mod tests {
275 use super::*;
276 use crate::storage::SledBackend;
277 use tempfile::TempDir;
278
279 fn create_test_inventory() -> Inventory {
280 let mut inventory = Inventory::new();
281
282 let mut node1 = TestNode::new("test_math.py::test_add", "test_math.py");
283 node1.add_marker("unit");
284 node1.build_keywords();
285 inventory.add(node1);
286
287 let mut node2 = TestNode::new("test_math.py::test_subtract", "test_math.py");
288 node2.add_marker("unit");
289 node2.build_keywords();
290 inventory.add(node2);
291
292 let mut node3 = TestNode::new("test_api.py::test_login", "test_api.py");
293 node3.add_marker("integration");
294 node3.add_marker("slow");
295 node3.build_keywords();
296 inventory.add(node3);
297
298 inventory
299 }
300
301 #[test]
302 fn test_basic_operations() {
303 let inventory = create_test_inventory();
304
305 assert_eq!(inventory.len(), 3);
306 assert!(inventory.contains("test_math.py::test_add"));
307 assert!(!inventory.contains("nonexistent"));
308
309 let node = inventory.get("test_api.py::test_login").unwrap();
310 assert!(node.has_marker("integration"));
311 }
312
313 #[test]
314 fn test_keyword_filter() {
315 let inventory = create_test_inventory();
316
317 let results = inventory.filter_by_keyword("math");
318 assert_eq!(results.len(), 2);
319
320 let results = inventory.filter_by_keyword("login");
321 assert_eq!(results.len(), 1);
322 }
323
324 #[test]
325 fn test_marker_filter() {
326 let inventory = create_test_inventory();
327
328 let results = inventory.filter_by_marker("unit");
329 assert_eq!(results.len(), 2);
330
331 let results = inventory.filter_by_marker("integration");
332 assert_eq!(results.len(), 1);
333
334 let results = inventory.filter_by_marker("slow");
335 assert_eq!(results.len(), 1);
336 }
337
338 #[test]
339 fn test_combined_filter() {
340 let inventory = create_test_inventory();
341
342 let results = inventory.filter(Some("test"), Some("unit"));
343 assert_eq!(results.len(), 2);
344
345 let results = inventory.filter(Some("login"), Some("integration"));
346 assert_eq!(results.len(), 1);
347
348 let results = inventory.filter(Some("math"), Some("integration"));
349 assert_eq!(results.len(), 0);
350 }
351
352 #[test]
353 fn test_persistence() {
354 let tmp = TempDir::new().unwrap();
355 let storage = SledBackend::open(tmp.path()).unwrap();
356
357 let mut inventory = create_test_inventory();
359 inventory.update_meta("abc123".to_string(), 1234567890);
360 inventory.save(&storage, "ctx-1").unwrap();
361
362 let loaded = Inventory::load(&storage, "ctx-1").unwrap();
364
365 assert_eq!(loaded.len(), 3);
366 assert_eq!(loaded.meta().hash, "abc123");
367 assert!(loaded.contains("test_math.py::test_add"));
368 }
369
370 #[test]
371 fn test_clear_storage() {
372 let tmp = TempDir::new().unwrap();
373 let storage = SledBackend::open(tmp.path()).unwrap();
374
375 let inventory = create_test_inventory();
376 inventory.save(&storage, "ctx-1").unwrap();
377
378 let loaded = Inventory::load(&storage, "ctx-1").unwrap();
380 assert_eq!(loaded.len(), 3);
381
382 Inventory::clear_storage(&storage, "ctx-1").unwrap();
384
385 let loaded = Inventory::load(&storage, "ctx-1").unwrap();
387 assert_eq!(loaded.len(), 0);
388 }
389}