Files
poimen-memory/tests/it_pg_repo.rs
T

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);
}