Files
clawhdf5/crates/clawhdf5-agent/tests/stress_tests.rs
T

691 lines
22 KiB
Rust

use clawhdf5_agent::bm25::BM25Index;
use clawhdf5_agent::hybrid::hybrid_search;
use clawhdf5_agent::vector_search;
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
use std::path::Path;
use tempfile::TempDir;
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
struct Rng(u32);
impl Rng {
fn new(seed: u32) -> Self {
Self(seed)
}
fn next_u32(&mut self) -> u32 {
self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345);
self.0 >> 16
}
fn next_f32(&mut self) -> f32 {
self.next_u32() as f32 / 65536.0 - 0.5
}
fn next_usize(&mut self, max: usize) -> usize {
self.next_u32() as usize % max
}
}
fn make_vec(rng: &mut Rng, dim: usize) -> Vec<f32> {
(0..dim).map(|_| rng.next_f32()).collect()
}
fn make_config(dir: &TempDir, dim: usize) -> MemoryConfig {
MemoryConfig::new(dir.path().join("test.h5"), "stress-agent", dim)
}
fn make_entry_simple(i: usize, dim: usize) -> MemoryEntry {
MemoryEntry {
chunk: format!("chunk_{i}"),
embedding: vec![i as f32 * 0.001; dim],
source_channel: "stress".into(),
timestamp: i as f64,
session_id: format!("sess_{}", i / 100),
tags: String::new(),
}
}
fn file_size(path: &Path) -> u64 {
std::fs::metadata(path).map(|m| m.len()).unwrap_or(0)
}
// ---------------------------------------------------------------------------
// 1. 100K entries in batches of 1000
// ---------------------------------------------------------------------------
#[test]
fn test_100k_entries() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 4);
let mut mem = HDF5Memory::create(config).unwrap();
for batch in 0..100 {
let start = batch * 1000;
let entries: Vec<MemoryEntry> = (start..start + 1000)
.map(|i| make_entry_simple(i, 4))
.collect();
mem.save_batch(entries).unwrap();
}
assert_eq!(mem.count(), 100_000);
assert_eq!(mem.count_active(), 100_000);
}
// ---------------------------------------------------------------------------
// 2. Heavy tombstoning: save 10K, delete 5K, compact, verify 5K remain
// ---------------------------------------------------------------------------
#[test]
fn test_heavy_tombstoning() {
let dir = TempDir::new().unwrap();
let mut config = make_config(&dir, 4);
config.compact_threshold = 0.0; // disable auto-compact
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
let entries: Vec<MemoryEntry> = (0..10_000).map(|i| make_entry_simple(i, 4)).collect();
mem.save_batch(entries).unwrap();
assert_eq!(mem.count(), 10_000);
// Delete every other entry (5000 deletions)
let mut rng = Rng::new(42);
let mut deleted = std::collections::HashSet::new();
while deleted.len() < 5000 {
let idx = rng.next_usize(10_000);
if deleted.insert(idx) {
mem.delete(idx).unwrap();
}
}
assert_eq!(mem.count_active(), 5000);
let removed = mem.compact().unwrap();
assert_eq!(removed, 5000);
assert_eq!(mem.count(), 5000);
assert_eq!(mem.count_active(), 5000);
// Verify persistence
let reopened = HDF5Memory::open(&path).unwrap();
assert_eq!(reopened.count(), 5000);
}
// ---------------------------------------------------------------------------
// 3. Repeated open/close cycles
// ---------------------------------------------------------------------------
#[test]
fn test_repeated_open_close() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 4);
let path = config.path.clone();
{
let mut mem = HDF5Memory::create(config).unwrap();
let entries: Vec<MemoryEntry> = (0..100).map(|i| make_entry_simple(i, 4)).collect();
mem.save_batch(entries).unwrap();
}
{
let mut mem = HDF5Memory::open(&path).unwrap();
assert_eq!(mem.count(), 100);
let entries: Vec<MemoryEntry> = (100..200).map(|i| make_entry_simple(i, 4)).collect();
mem.save_batch(entries).unwrap();
}
let mem = HDF5Memory::open(&path).unwrap();
assert_eq!(mem.count(), 200);
assert_eq!(mem.count_active(), 200);
}
// ---------------------------------------------------------------------------
// 4. Large embeddings (1536-dim, ada-002 size) x 10K
// ---------------------------------------------------------------------------
#[test]
fn test_large_embeddings_1536() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 1536);
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
let mut rng = Rng::new(42);
let entries: Vec<MemoryEntry> = (0..10_000)
.map(|i| MemoryEntry {
chunk: format!("large_emb_{i}"),
embedding: make_vec(&mut rng, 1536),
source_channel: "test".into(),
timestamp: i as f64,
session_id: "s1".into(),
tags: String::new(),
})
.collect();
mem.save_batch(entries).unwrap();
assert_eq!(mem.count(), 10_000);
// Verify persistence
let reopened = HDF5Memory::open(&path).unwrap();
assert_eq!(reopened.count(), 10_000);
// Verify search works on large dims
let (_, cache, _, _) = clawhdf5_agent::storage::read_from_disk(&path).unwrap();
let query = make_vec(&mut Rng::new(99), 1536);
let results =
vector_search::cosine_similarity_batch(&query, &cache.embeddings, &cache.tombstones);
assert_eq!(results.len(), 10_000);
}
// ---------------------------------------------------------------------------
// 5. Concurrent-like pattern: interleaved save, search, delete, compact
// ---------------------------------------------------------------------------
#[test]
fn test_concurrent_like_pattern() {
let dir = TempDir::new().unwrap();
let mut config = make_config(&dir, 8);
config.compact_threshold = 0.0;
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
let mut rng = Rng::new(42);
// Round 1: save 100 entries
for i in 0..100 {
mem.save(MemoryEntry {
chunk: format!("round1_{i}"),
embedding: make_vec(&mut rng, 8),
source_channel: "test".into(),
timestamp: i as f64,
session_id: "s1".into(),
tags: String::new(),
})
.unwrap();
}
assert_eq!(mem.count(), 100);
// Flush WAL so read_from_disk sees the data
mem.flush_wal().unwrap();
// Read cache for search
let (_, cache, _, _) = clawhdf5_agent::storage::read_from_disk(&path).unwrap();
let query = make_vec(&mut Rng::new(99), 8);
let results =
vector_search::cosine_similarity_batch(&query, &cache.embeddings, &cache.tombstones);
assert_eq!(results.len(), 100);
// Delete 20 entries
for i in 0..20 {
mem.delete(i).unwrap();
}
assert_eq!(mem.count_active(), 80);
// Compact
let removed = mem.compact().unwrap();
assert_eq!(removed, 20);
// Round 2: save 50 more
for i in 0..50 {
mem.save(MemoryEntry {
chunk: format!("round2_{i}"),
embedding: make_vec(&mut rng, 8),
source_channel: "test".into(),
timestamp: 200.0 + i as f64,
session_id: "s2".into(),
tags: String::new(),
})
.unwrap();
}
assert_eq!(mem.count_active(), 130);
// Flush WAL so read_from_disk sees round 2 entries
mem.flush_wal().unwrap();
// Final search
let (_, cache2, _, _) = clawhdf5_agent::storage::read_from_disk(&path).unwrap();
let results2 =
vector_search::cosine_similarity_batch(&query, &cache2.embeddings, &cache2.tombstones);
assert_eq!(results2.len(), 130);
}
// ---------------------------------------------------------------------------
// 6. File size growth is reasonable
// ---------------------------------------------------------------------------
#[test]
fn test_file_size_growth() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 8);
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
let mut sizes: Vec<u64> = Vec::new();
for batch in 0..10 {
let start = batch * 500;
let entries: Vec<MemoryEntry> = (start..start + 500)
.map(|i| make_entry_simple(i, 8))
.collect();
mem.save_batch(entries).unwrap();
sizes.push(file_size(&path));
}
// File size should grow roughly linearly (not exponentially)
// Check that doubling entries doesn't more than triple file size
for i in 1..sizes.len() {
assert!(
sizes[i] > sizes[i - 1],
"file should grow: sizes[{i}]={} <= sizes[{}]={}",
sizes[i],
i - 1,
sizes[i - 1]
);
}
// The 10x data file should be less than 15x the 1x data file
let ratio = sizes[9] as f64 / sizes[0] as f64;
assert!(
ratio < 15.0,
"file growth ratio too high: {ratio:.1}x for 10x data"
);
}
// ---------------------------------------------------------------------------
// 7. Vector search accuracy with known vectors
// ---------------------------------------------------------------------------
#[test]
fn test_vector_search_accuracy() {
let vectors = vec![
vec![1.0, 0.0, 0.0, 0.0], // idx 0: unit x
vec![0.0, 1.0, 0.0, 0.0], // idx 1: unit y
vec![0.0, 0.0, 1.0, 0.0], // idx 2: unit z
vec![1.0 / 2.0_f32.sqrt(), 1.0 / 2.0_f32.sqrt(), 0.0, 0.0], // idx 3: 45 deg
vec![-1.0, 0.0, 0.0, 0.0], // idx 4: opposite x
];
let tombstones = vec![0u8; 5];
let query = vec![1.0, 0.0, 0.0, 0.0]; // unit x
let results = vector_search::cosine_similarity_batch(&query, &vectors, &tombstones);
// Expected order: idx0 (1.0) > idx3 (~0.707) > idx1 (0.0) = idx2 (0.0) > idx4 (-1.0)
assert_eq!(results[0].0, 0);
assert!((results[0].1 - 1.0).abs() < 1e-5);
assert_eq!(results[1].0, 3);
assert!((results[1].1 - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-5);
// idx4 should be last with cos = -1.0
assert_eq!(results[4].0, 4);
assert!((results[4].1 - (-1.0)).abs() < 1e-5);
// Verify cosine_similarity standalone
let sim = vector_search::cosine_similarity(&query, &vectors[3]);
assert!((sim - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-5);
}
// ---------------------------------------------------------------------------
// 8. BM25 accuracy with known term frequencies
// ---------------------------------------------------------------------------
#[test]
fn test_bm25_accuracy() {
let docs = vec![
"rust rust rust systems programming".to_string(), // idx 0: 3x "rust"
"rust programming language".to_string(), // idx 1: 1x "rust"
"python scripting language".to_string(), // idx 2: 0x "rust"
"java enterprise rust system".to_string(), // idx 3: 1x "rust"
"javascript web development frontend".to_string(), // idx 4: 0x "rust"
];
let tombstones = vec![0u8; 5];
let index = BM25Index::build(&docs, &tombstones);
// Query for "rust"
let results = index.search("rust", 10);
// Doc 0 should rank first (highest TF for "rust")
assert!(!results.is_empty());
assert_eq!(results[0].0, 0, "doc with 3x 'rust' should rank first");
// Docs 2 and 4 should not appear (no "rust")
let result_ids: Vec<usize> = results.iter().map(|(id, _)| *id).collect();
assert!(
!result_ids.contains(&2),
"doc without 'rust' should not appear"
);
assert!(
!result_ids.contains(&4),
"doc without 'rust' should not appear"
);
// Query for rare term "enterprise"
let rare_results = index.search("enterprise", 10);
assert_eq!(rare_results.len(), 1);
assert_eq!(rare_results[0].0, 3);
// Multi-term query: "rust programming" should boost doc 0 and 1
let multi = index.search("rust programming", 10);
assert!(multi.len() >= 2);
let top2: Vec<usize> = multi.iter().take(2).map(|(id, _)| *id).collect();
assert!(
top2.contains(&0),
"doc 0 should be in top 2 for 'rust programming'"
);
assert!(
top2.contains(&1),
"doc 1 should be in top 2 for 'rust programming'"
);
}
// ---------------------------------------------------------------------------
// 9. Hybrid search correctness: vector and keyword disagree
// ---------------------------------------------------------------------------
#[test]
fn test_hybrid_search_correctness() {
// Doc 0: great vector match, no keyword match
// Doc 1: no vector match, great keyword match
// Doc 2: moderate vector match, moderate keyword match
// Doc 3: some vector, some keyword
// Doc 4: filler
let vectors = vec![
vec![1.0, 0.0, 0.0, 0.0], // idx 0: identical to query
vec![0.0, 1.0, 0.0, 0.0], // idx 1: orthogonal
vec![0.7, 0.7, 0.0, 0.0], // idx 2: partial match
vec![0.3, 0.3, 0.3, 0.3], // idx 3: mild match
vec![0.0, 0.0, 0.0, 1.0], // idx 4: orthogonal
];
let chunks = vec![
"gamma delta epsilon phi".to_string(), // 0: no keyword match
"alpha alpha alpha beta alpha".to_string(), // 1: heavy keyword
"alpha gamma delta".to_string(), // 2: some keyword
"alpha beta gamma".to_string(), // 3: some keyword
"alpha omega sigma".to_string(), // 4: some keyword
];
let tombstones = vec![0u8; 5];
let bm25 = BM25Index::build(&chunks, &tombstones);
let query_emb = vec![1.0, 0.0, 0.0, 0.0];
// Vector-only: doc 0 should win
let vec_only = hybrid_search(
&query_emb,
"alpha",
&vectors,
&chunks,
&tombstones,
&bm25,
1.0,
0.0,
5,
);
assert_eq!(vec_only[0].0, 0, "vector-only: doc 0 should win");
// Keyword-only: doc 1 should win (most "alpha" occurrences)
let kw_only = hybrid_search(
&query_emb,
"alpha",
&vectors,
&chunks,
&tombstones,
&bm25,
0.0,
1.0,
5,
);
assert_eq!(kw_only[0].0, 1, "keyword-only: doc 1 should win");
// Balanced: both doc 0 and doc 1 should appear in top 3
let balanced = hybrid_search(
&query_emb,
"alpha",
&vectors,
&chunks,
&tombstones,
&bm25,
0.5,
0.5,
5,
);
let top3: Vec<usize> = balanced.iter().take(3).map(|(id, _)| *id).collect();
assert!(top3.contains(&0), "balanced: doc 0 should be in top 3");
assert!(top3.contains(&1), "balanced: doc 1 should be in top 3");
}
// ---------------------------------------------------------------------------
// 10. Session tracking stress: 1000 sessions
// ---------------------------------------------------------------------------
#[test]
fn test_session_tracking_stress() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 4);
let path = config.path.clone();
{
let mut mem = HDF5Memory::create(config).unwrap();
for i in 0..1000 {
mem.add_session(
&format!("sess_{i}"),
i * 10,
(i + 1) * 10,
"api",
&format!("Summary for session {i}"),
)
.unwrap();
}
}
// Reopen and verify random sessions
let mem = HDF5Memory::open(&path).unwrap();
let mut rng = Rng::new(42);
for _ in 0..100 {
let idx = rng.next_usize(1000);
let summary = mem
.get_session_summary(&format!("sess_{idx}"))
.unwrap()
.unwrap();
assert_eq!(summary, format!("Summary for session {idx}"));
}
// Non-existent session returns None
assert!(mem.get_session_summary("nonexistent").unwrap().is_none());
}
// ---------------------------------------------------------------------------
// 11. Varying batch sizes
// ---------------------------------------------------------------------------
#[test]
fn test_batch_sizes_vary() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 4);
let mut mem = HDF5Memory::create(config).unwrap();
// Batch of 1
mem.save_batch(vec![make_entry_simple(0, 4)]).unwrap();
assert_eq!(mem.count(), 1);
// Batch of 10
let batch10: Vec<MemoryEntry> = (1..11).map(|i| make_entry_simple(i, 4)).collect();
mem.save_batch(batch10).unwrap();
assert_eq!(mem.count(), 11);
// Batch of 500
let batch500: Vec<MemoryEntry> = (11..511).map(|i| make_entry_simple(i, 4)).collect();
mem.save_batch(batch500).unwrap();
assert_eq!(mem.count(), 511);
// Single saves
for i in 511..521 {
mem.save(make_entry_simple(i, 4)).unwrap();
}
assert_eq!(mem.count(), 521);
assert_eq!(mem.count_active(), 521);
}
// ---------------------------------------------------------------------------
// 12. Delete all entries
// ---------------------------------------------------------------------------
#[test]
fn test_delete_all_entries() {
let dir = TempDir::new().unwrap();
let mut config = make_config(&dir, 4);
config.compact_threshold = 0.0;
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
let entries: Vec<MemoryEntry> = (0..100).map(|i| make_entry_simple(i, 4)).collect();
mem.save_batch(entries).unwrap();
for i in 0..100 {
mem.delete(i).unwrap();
}
assert_eq!(mem.count(), 100);
assert_eq!(mem.count_active(), 0);
let removed = mem.compact().unwrap();
assert_eq!(removed, 100);
assert_eq!(mem.count(), 0);
// Verify persistence
let reopened = HDF5Memory::open(&path).unwrap();
assert_eq!(reopened.count(), 0);
}
// ---------------------------------------------------------------------------
// 13. Compact empty file
// ---------------------------------------------------------------------------
#[test]
fn test_compact_empty() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 4);
let mut mem = HDF5Memory::create(config).unwrap();
let removed = mem.compact().unwrap();
assert_eq!(removed, 0);
assert_eq!(mem.count(), 0);
}
// ---------------------------------------------------------------------------
// 14. Search with all entries tombstoned
// ---------------------------------------------------------------------------
#[test]
fn test_search_all_tombstoned() {
let vectors = vec![
vec![1.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0],
vec![0.0, 0.0, 1.0],
];
let tombstones = vec![1u8; 3]; // all tombstoned
let query = vec![1.0, 0.0, 0.0];
let vec_results = vector_search::cosine_similarity_batch(&query, &vectors, &tombstones);
assert!(vec_results.is_empty());
let chunks = vec!["hello".to_string(), "world".to_string(), "foo".to_string()];
let bm25 = BM25Index::build(&chunks, &tombstones);
let bm25_results = bm25.search("hello", 10);
assert!(bm25_results.is_empty());
let hybrid_results = hybrid_search(
&query,
"hello",
&vectors,
&chunks,
&tombstones,
&bm25,
0.5,
0.5,
10,
);
assert!(hybrid_results.is_empty());
}
// ---------------------------------------------------------------------------
// 15. Unicode content handling
// ---------------------------------------------------------------------------
#[test]
fn test_unicode_content() {
let dir = TempDir::new().unwrap();
let config = make_config(&dir, 4);
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
let entries = vec![
MemoryEntry {
chunk: "Hello world in Japanese: \u{3053}\u{3093}\u{306b}\u{3061}\u{306f}".into(),
embedding: vec![1.0, 0.0, 0.0, 0.0],
source_channel: "test".into(),
timestamp: 1.0,
session_id: "s1".into(),
tags: "\u{00e9}m\u{00f6}ji".into(),
},
MemoryEntry {
chunk: "Chinese: \u{4f60}\u{597d}\u{4e16}\u{754c}".into(),
embedding: vec![0.0, 1.0, 0.0, 0.0],
source_channel: "test".into(),
timestamp: 2.0,
session_id: "s1".into(),
tags: String::new(),
},
MemoryEntry {
chunk: "Emoji test: \u{1f600}\u{1f680}\u{2764}".into(),
embedding: vec![0.0, 0.0, 1.0, 0.0],
source_channel: "test".into(),
timestamp: 3.0,
session_id: "s1".into(),
tags: String::new(),
},
];
mem.save_batch(entries).unwrap();
let reopened = HDF5Memory::open(&path).unwrap();
assert_eq!(reopened.count(), 3);
let (_, cache, _, _) = clawhdf5_agent::storage::read_from_disk(&path).unwrap();
assert!(cache.chunks[0].contains("\u{3053}\u{3093}\u{306b}\u{3061}\u{306f}"));
assert!(cache.chunks[1].contains("\u{4f60}\u{597d}"));
}
// ---------------------------------------------------------------------------
// 16. Rapid save/delete cycles
// ---------------------------------------------------------------------------
#[test]
fn test_rapid_save_delete_cycles() {
let dir = TempDir::new().unwrap();
let mut config = make_config(&dir, 4);
config.compact_threshold = 0.0;
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
// 50 cycles of: save 10, delete oldest 5
let mut next_id = 0usize;
let mut active_start = 0usize;
for _ in 0..50 {
let entries: Vec<MemoryEntry> = (next_id..next_id + 10)
.map(|i| make_entry_simple(i, 4))
.collect();
mem.save_batch(entries).unwrap();
next_id += 10;
for i in active_start..active_start + 5 {
mem.delete(i).unwrap();
}
active_start += 5;
}
// We saved 500 entries total, deleted 250
assert_eq!(mem.count(), 500);
assert_eq!(mem.count_active(), 250);
// Compact and verify
let removed = mem.compact().unwrap();
assert_eq!(removed, 250);
assert_eq!(mem.count(), 250);
let reopened = HDF5Memory::open(&path).unwrap();
assert_eq!(reopened.count(), 250);
}