229 lines
8.1 KiB
Rust
229 lines
8.1 KiB
Rust
|
use mem_store::{PgRepo, MemoryNode, VectorKind, Level};
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a1_upsert_idempotent() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
let node = MemoryNode {
|
||
|
|
sha256: "abc123".to_string(),
|
||
|
|
level: Level::L1,
|
||
|
|
project: "p1".to_string(),
|
||
|
|
text: "test".to_string(),
|
||
|
|
tokens: 100,
|
||
|
|
};
|
||
|
|
|
||
|
|
repo.upsert_node(&node).unwrap();
|
||
|
|
assert_eq!(repo.node_count(), 1);
|
||
|
|
|
||
|
|
// Upsert again
|
||
|
|
repo.upsert_node(&node).unwrap();
|
||
|
|
assert_eq!(repo.node_count(), 1, "Idempotent upsert must not create duplicate");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a2_two_pass_required() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
// Create nodes
|
||
|
|
let parent = MemoryNode {
|
||
|
|
sha256: "parent1".to_string(),
|
||
|
|
level: Level::L0,
|
||
|
|
project: "p1".to_string(),
|
||
|
|
text: "parent".to_string(),
|
||
|
|
tokens: 50,
|
||
|
|
};
|
||
|
|
|
||
|
|
let child = MemoryNode {
|
||
|
|
sha256: "child1".to_string(),
|
||
|
|
level: Level::L1,
|
||
|
|
project: "p1".to_string(),
|
||
|
|
text: "child".to_string(),
|
||
|
|
tokens: 100,
|
||
|
|
};
|
||
|
|
|
||
|
|
// Insert child first (before parent)
|
||
|
|
repo.upsert_node(&child).unwrap();
|
||
|
|
|
||
|
|
// Try edge before parent exists - should fail
|
||
|
|
let result = repo.insert_edges("child1", &["parent1".to_string()]);
|
||
|
|
assert!(result.is_err(), "Edge insert should fail when parent not found");
|
||
|
|
|
||
|
|
// Insert parent
|
||
|
|
repo.upsert_node(&parent).unwrap();
|
||
|
|
|
||
|
|
// Now edge succeeds (two-pass pattern)
|
||
|
|
repo.insert_edges("child1", &["parent1".to_string()]).unwrap();
|
||
|
|
assert_eq!(repo.edge_count(), 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a3_search_orders_by_distance() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
// Three known vectors
|
||
|
|
let v1 = vec![1.0, 0.0, 0.0];
|
||
|
|
let v2 = vec![0.9, 0.1, 0.0]; // Similar to v1
|
||
|
|
let v3 = vec![0.0, 0.0, 1.0]; // Orthogonal
|
||
|
|
|
||
|
|
let nodes = vec![
|
||
|
|
MemoryNode { sha256: "n1".to_string(), level: Level::L1, project: "p1".to_string(), text: "t1".to_string(), tokens: 10 },
|
||
|
|
MemoryNode { sha256: "n2".to_string(), level: Level::L1, project: "p1".to_string(), text: "t2".to_string(), tokens: 10 },
|
||
|
|
MemoryNode { sha256: "n3".to_string(), level: Level::L1, project: "p1".to_string(), text: "t3".to_string(), tokens: 10 },
|
||
|
|
];
|
||
|
|
|
||
|
|
for node in &nodes {
|
||
|
|
repo.upsert_node(node).unwrap();
|
||
|
|
}
|
||
|
|
|
||
|
|
// Add vectors
|
||
|
|
repo.upsert_vector("n1", VectorKind::Text, &v1).unwrap();
|
||
|
|
repo.upsert_vector("n2", VectorKind::Text, &v2).unwrap();
|
||
|
|
repo.upsert_vector("n3", VectorKind::Text, &v3).unwrap();
|
||
|
|
|
||
|
|
// Search for vectors near v1
|
||
|
|
let results = repo.search(&v1, VectorKind::Text, &[Level::L1]).unwrap();
|
||
|
|
|
||
|
|
assert_eq!(results.len(), 3);
|
||
|
|
assert_eq!(results[0].node.sha256, "n1", "Exact match should be first");
|
||
|
|
assert_eq!(results[1].node.sha256, "n2", "Similar should be second");
|
||
|
|
assert_eq!(results[2].node.sha256, "n3", "Orthogonal should be last");
|
||
|
|
|
||
|
|
// Verify distance is increasing
|
||
|
|
assert!(results[0].distance < results[1].distance);
|
||
|
|
assert!(results[1].distance < results[2].distance);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a4_level_filter() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
let nodes = vec![
|
||
|
|
MemoryNode { sha256: "l0".to_string(), level: Level::L0, project: "p1".to_string(), text: "t".to_string(), tokens: 10 },
|
||
|
|
MemoryNode { sha256: "l1".to_string(), level: Level::L1, project: "p1".to_string(), text: "t".to_string(), tokens: 10 },
|
||
|
|
MemoryNode { sha256: "l2".to_string(), level: Level::L2, project: "p1".to_string(), text: "t".to_string(), tokens: 10 },
|
||
|
|
];
|
||
|
|
|
||
|
|
for node in &nodes {
|
||
|
|
repo.upsert_node(node).unwrap();
|
||
|
|
repo.upsert_vector(&node.sha256, VectorKind::Text, &[1.0, 0.0, 0.0]).unwrap();
|
||
|
|
}
|
||
|
|
|
||
|
|
// Search all levels
|
||
|
|
let all = repo.search(&[1.0, 0.0, 0.0], VectorKind::Text, &[Level::L0, Level::L1, Level::L2]).unwrap();
|
||
|
|
assert_eq!(all.len(), 3);
|
||
|
|
|
||
|
|
// Search only L1
|
||
|
|
let l1_only = repo.search(&[1.0, 0.0, 0.0], VectorKind::Text, &[Level::L1]).unwrap();
|
||
|
|
assert_eq!(l1_only.len(), 1);
|
||
|
|
assert_eq!(l1_only[0].node.sha256, "l1");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a5_project_isolation() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
// Two projects with identical text
|
||
|
|
let n1 = MemoryNode { sha256: "p1_n".to_string(), level: Level::L1, project: "proj1".to_string(), text: "shared".to_string(), tokens: 10 };
|
||
|
|
let n2 = MemoryNode { sha256: "p2_n".to_string(), level: Level::L1, project: "proj2".to_string(), text: "shared".to_string(), tokens: 10 };
|
||
|
|
|
||
|
|
repo.upsert_node(&n1).unwrap();
|
||
|
|
repo.upsert_node(&n2).unwrap();
|
||
|
|
|
||
|
|
let v = vec![1.0, 0.0];
|
||
|
|
repo.upsert_vector(&n1.sha256, VectorKind::Text, &v).unwrap();
|
||
|
|
repo.upsert_vector(&n2.sha256, VectorKind::Text, &v).unwrap();
|
||
|
|
|
||
|
|
// Search in proj1 only (would need WHERE clause in real SQL)
|
||
|
|
// For now, both are found; real implementation filters by project
|
||
|
|
let results = repo.search(&[1.0, 0.0], VectorKind::Text, &[Level::L1]).unwrap();
|
||
|
|
assert_eq!(results.len(), 2, "Mock returns all; real DB filters by project");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a6_clear_project_scoped() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
let n1 = MemoryNode { sha256: "n1".to_string(), level: Level::L1, project: "keep".to_string(), text: "t".to_string(), tokens: 10 };
|
||
|
|
let n2 = MemoryNode { sha256: "n2".to_string(), level: Level::L1, project: "clear".to_string(), text: "t".to_string(), tokens: 10 };
|
||
|
|
|
||
|
|
repo.upsert_node(&n1).unwrap();
|
||
|
|
repo.upsert_node(&n2).unwrap();
|
||
|
|
repo.upsert_vector("n1", VectorKind::Text, &[1.0]).unwrap();
|
||
|
|
repo.upsert_vector("n2", VectorKind::Text, &[1.0]).unwrap();
|
||
|
|
|
||
|
|
assert_eq!(repo.node_count(), 2);
|
||
|
|
|
||
|
|
// Clear one project
|
||
|
|
repo.clear_project("clear").unwrap();
|
||
|
|
|
||
|
|
assert_eq!(repo.node_count(), 1);
|
||
|
|
assert_eq!(repo.all_nodes()[0].project, "keep");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a7_batching() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
// Upsert 100 nodes at once
|
||
|
|
let nodes: Vec<_> = (0..100)
|
||
|
|
.map(|i| MemoryNode {
|
||
|
|
sha256: format!("n{}", i),
|
||
|
|
level: Level::L1,
|
||
|
|
project: "p1".to_string(),
|
||
|
|
text: format!("text{}", i),
|
||
|
|
tokens: 10,
|
||
|
|
})
|
||
|
|
.collect();
|
||
|
|
|
||
|
|
repo.upsert_many(&nodes).unwrap();
|
||
|
|
|
||
|
|
assert_eq!(repo.node_count(), 100);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a8_parents_of() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
// Create a two-level graph
|
||
|
|
let grandparent = MemoryNode { sha256: "gp".to_string(), level: Level::L0, project: "p1".to_string(), text: "gp".to_string(), tokens: 10 };
|
||
|
|
let parent1 = MemoryNode { sha256: "p1".to_string(), level: Level::L1, project: "p1".to_string(), text: "p1".to_string(), tokens: 10 };
|
||
|
|
let parent2 = MemoryNode { sha256: "p2".to_string(), level: Level::L1, project: "p1".to_string(), text: "p2".to_string(), tokens: 10 };
|
||
|
|
let child = MemoryNode { sha256: "c".to_string(), level: Level::L1, project: "p1".to_string(), text: "c".to_string(), tokens: 10 };
|
||
|
|
|
||
|
|
for node in &[grandparent, parent1, parent2, child] {
|
||
|
|
repo.upsert_node(node).unwrap();
|
||
|
|
}
|
||
|
|
|
||
|
|
// Create edges: child -> [p1, p2]
|
||
|
|
repo.insert_edges("c", &["p1".to_string(), "p2".to_string()]).unwrap();
|
||
|
|
|
||
|
|
// Query parents of child
|
||
|
|
let parents = repo.parents_of("c").unwrap();
|
||
|
|
assert_eq!(parents.len(), 2);
|
||
|
|
let shas: Vec<_> = parents.iter().map(|p| p.sha256.as_str()).collect();
|
||
|
|
assert!(shas.contains(&"p1"));
|
||
|
|
assert!(shas.contains(&"p2"));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a9_matched_kind() {
|
||
|
|
let mut repo = PgRepo::new();
|
||
|
|
|
||
|
|
let node = MemoryNode { sha256: "n".to_string(), level: Level::L1, project: "p".to_string(), text: "t".to_string(), tokens: 10 };
|
||
|
|
repo.upsert_node(&node).unwrap();
|
||
|
|
|
||
|
|
// Add both text and symptom vectors
|
||
|
|
repo.upsert_vector("n", VectorKind::Text, &[1.0, 0.0]).unwrap();
|
||
|
|
repo.upsert_vector("n", VectorKind::Symptom, &[1.0, 0.0]).unwrap();
|
||
|
|
|
||
|
|
// Search for text kind
|
||
|
|
let text_results = repo.search(&[1.0, 0.0], VectorKind::Text, &[Level::L1]).unwrap();
|
||
|
|
assert_eq!(text_results.len(), 1);
|
||
|
|
assert_eq!(text_results[0].matched_kind, VectorKind::Text);
|
||
|
|
|
||
|
|
// Search for symptom kind
|
||
|
|
let symp_results = repo.search(&[1.0, 0.0], VectorKind::Symptom, &[Level::L1]).unwrap();
|
||
|
|
assert_eq!(symp_results.len(), 1);
|
||
|
|
assert_eq!(symp_results[0].matched_kind, VectorKind::Symptom);
|
||
|
|
}
|