use edgestore::{
text::types::{FacetValue, TextRecord},
EdgestoreConfig, Engine, TextEngine, VectorEngine,
};
use tempfile::TempDir;
fn open_engine(dir: &TempDir) -> Engine {
Engine::open(EdgestoreConfig::new(dir.path())).unwrap()
}
#[test]
fn test_index_and_search_basic() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"The quick brown fox",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc2",
"The lazy dog sleeps",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc3",
"Quick brown fox jumps",
std::collections::HashMap::new(),
)
.unwrap();
let results = engine.search_text(b"ns", "quick brown", 3).unwrap();
assert!(!results.is_empty(), "search should return results");
assert!(
results.iter().any(|r| r.doc_id == b"doc1"),
"doc1 should match 'quick brown'"
);
assert!(
results.iter().any(|r| r.doc_id == b"doc3"),
"doc3 should match 'quick brown'"
);
}
#[test]
fn test_bm25_ranking() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc2",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
let results = engine.search_text(b"ns", "hello", 2).unwrap();
assert_eq!(results.len(), 2);
assert_eq!(
results[0].doc_id, b"doc1",
"doc with more 'hello' should rank higher"
);
assert!(results[0].score > results[1].score);
}
#[test]
fn test_search_empty_namespace() {
let dir = TempDir::new().unwrap();
let engine = open_engine(&dir);
let results = engine.search_text(b"ns", "hello", 5).unwrap();
assert!(results.is_empty());
}
#[test]
fn test_search_empty_query() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
let results = engine.search_text(b"ns", "", 5).unwrap();
assert!(results.is_empty());
let results2 = engine.search_text(b"ns", "the a an", 5).unwrap();
assert!(
results2.is_empty(),
"stopwords-only query should return empty"
);
}
#[test]
fn test_delete_removes_from_search() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
let results_before = engine.search_text(b"ns", "hello", 5).unwrap();
assert_eq!(results_before.len(), 1);
engine.delete_text(b"ns", b"doc1").unwrap();
let results_after = engine.search_text(b"ns", "hello", 5).unwrap();
assert!(
results_after.is_empty(),
"deleted doc should not appear in search"
);
}
#[test]
fn test_facet_filter() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
let mut facets1 = std::collections::HashMap::new();
facets1.insert(
"category".to_string(),
FacetValue::String("news".to_string()),
);
engine
.index_text(b"ns", b"doc1", "breaking news today", facets1)
.unwrap();
let mut facets2 = std::collections::HashMap::new();
facets2.insert(
"category".to_string(),
FacetValue::String("sports".to_string()),
);
engine
.index_text(b"ns", b"doc2", "sports update", facets2)
.unwrap();
let results = engine.search_text(b"ns", "news", 5).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].doc_id, b"doc1");
}
#[test]
fn test_search_ranking_stability() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"alpha beta gamma",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc2",
"beta gamma delta",
std::collections::HashMap::new(),
)
.unwrap();
let results1 = engine.search_text(b"ns", "beta gamma", 5).unwrap();
let results2 = engine.search_text(b"ns", "beta gamma", 5).unwrap();
assert_eq!(results1.len(), results2.len());
for (a, b) in results1.iter().zip(results2.iter()) {
assert_eq!(a.doc_id, b.doc_id);
assert!((a.score - b.score).abs() < 1e-6);
}
}
#[test]
fn test_index_text_record_retrieval() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
let mut facets = std::collections::HashMap::new();
facets.insert(
"author".to_string(),
FacetValue::String("Alice".to_string()),
);
facets.insert("views".to_string(), FacetValue::Number(42));
facets.insert("published".to_string(), FacetValue::Bool(true));
engine
.index_text(b"ns", b"doc1", "hello world", facets.clone())
.unwrap();
let text_ns = edgestore::text_namespace(b"ns");
let raw = engine.get(&text_ns, b"doc1").unwrap().unwrap();
let record = edgestore::decode_text_record(&raw).unwrap();
assert_eq!(record.text, "hello world");
assert_eq!(
record.facets.get("author"),
Some(&FacetValue::String("Alice".to_string()))
);
assert_eq!(record.facets.get("views"), Some(&FacetValue::Number(42)));
assert_eq!(
record.facets.get("published"),
Some(&FacetValue::Bool(true))
);
}
#[test]
fn test_reindex_updates_merged_index() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
let results = engine.search_text(b"ns", "hello", 5).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].doc_id, b"doc1");
engine
.index_text(b"ns", b"doc1", "foo bar", std::collections::HashMap::new())
.unwrap();
let results_hello = engine.search_text(b"ns", "hello", 5).unwrap();
assert!(
results_hello.is_empty(),
"old term 'hello' should not find re-indexed doc"
);
let results_foo = engine.search_text(b"ns", "foo", 5).unwrap();
assert_eq!(results_foo.len(), 1);
assert_eq!(results_foo[0].doc_id, b"doc1");
}
#[test]
fn test_incremental_index_many_docs() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
for i in 0..100 {
let text = format!("document number {} contains quick brown fox", i);
let key = format!("doc{:04}", i);
engine
.index_text(
b"ns",
key.as_bytes(),
&text,
std::collections::HashMap::new(),
)
.unwrap();
}
let results = engine.search_text(b"ns", "quick brown", 200).unwrap();
assert_eq!(
results.len(),
100,
"all 100 docs should match 'quick brown'"
);
for i in (0..100).step_by(2) {
let key = format!("doc{:04}", i);
engine.delete_text(b"ns", key.as_bytes()).unwrap();
}
let results_after = engine.search_text(b"ns", "quick brown", 200).unwrap();
assert_eq!(
results_after.len(),
50,
"50 docs should remain after deletion"
);
}
#[test]
fn test_namespace_isolation() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns1",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(b"ns2", b"doc1", "foo bar", std::collections::HashMap::new())
.unwrap();
let results1 = engine.search_text(b"ns1", "hello", 5).unwrap();
assert_eq!(results1.len(), 1);
let results2 = engine.search_text(b"ns2", "hello", 5).unwrap();
assert!(results2.is_empty(), "ns2 should not find ns1 terms");
let results3 = engine.search_text(b"ns2", "foo", 5).unwrap();
assert_eq!(results3.len(), 1);
}
#[test]
fn test_delete_all_docs_removes_index() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc2",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine.delete_text(b"ns", b"doc1").unwrap();
engine.delete_text(b"ns", b"doc2").unwrap();
let text_ns = edgestore::text_namespace(b"ns");
let index_bytes = engine.get(&text_ns, b"__index__").unwrap();
assert!(
index_bytes.is_none(),
"merged index should be deleted when all docs removed"
);
}
#[test]
fn test_search_performance_at_scale() {
use std::time::Instant;
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
let n = 10_000;
for i in 0..n {
let text = format!(
"document number {} contains quick brown fox jumps over lazy dog",
i
);
let key = format!("doc{:08}", i);
engine
.index_text(
b"ns",
key.as_bytes(),
&text,
std::collections::HashMap::new(),
)
.unwrap();
}
let start = Instant::now();
for _ in 0..100 {
let results = engine.search_text(b"ns", "quick brown fox", 10).unwrap();
assert!(!results.is_empty());
}
let elapsed = start.elapsed();
let avg_us = elapsed.as_micros() as f64 / 100.0;
let threshold_us = if cfg!(debug_assertions) {
50_000.0
} else {
5_000.0
};
assert!(
avg_us < threshold_us,
"search too slow: {:.1} µs at {} docs (threshold: {:.0} µs)",
avg_us,
n,
threshold_us
);
}
#[test]
fn test_cold_cache_search() {
let dir = TempDir::new().unwrap();
{
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc2",
"hello foo",
std::collections::HashMap::new(),
)
.unwrap();
engine.flush().unwrap();
}
{
let engine = open_engine(&dir);
let results = engine.search_text(b"ns", "hello", 5).unwrap();
assert_eq!(
results.len(),
2,
"cold-cache search should find both docs via disk fallback"
);
}
}
#[test]
fn test_typo_tolerance() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine
.index_text(
b"ns",
b"doc2",
"helo there",
std::collections::HashMap::new(),
)
.unwrap();
let exact = engine
.search_text_with_options(
b"ns",
"hello",
&edgestore::SearchOptions {
k: 5,
typo_tolerance: true,
..Default::default()
},
)
.unwrap();
assert!(
exact.iter().any(|r| r.doc_id == b"doc1"),
"exact match doc1 should be found"
);
assert!(
exact.iter().any(|r| r.doc_id == b"doc2"),
"typo-tolerant match doc2 ('helo' ~ 'hello') should be found"
);
}
#[test]
fn test_delete_fallback_cache_miss() {
let dir = TempDir::new().unwrap();
{
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine.flush().unwrap();
}
{
let mut engine = open_engine(&dir);
let results_before = engine.search_text(b"ns", "hello", 5).unwrap();
assert_eq!(results_before.len(), 1);
engine.delete_text(b"ns", b"doc1").unwrap();
let results_after = engine.search_text(b"ns", "hello", 5).unwrap();
assert!(
results_after.is_empty(),
"delete from cold cache should remove doc"
);
}
}
#[test]
fn test_reindex_with_facets() {
let dir = TempDir::new().unwrap();
let mut engine = open_engine(&dir);
let mut facets1 = std::collections::HashMap::new();
facets1.insert(
"category".to_string(),
FacetValue::String("news".to_string()),
);
engine
.index_text(b"ns", b"doc1", "breaking news today", facets1)
.unwrap();
let mut facets2 = std::collections::HashMap::new();
facets2.insert(
"category".to_string(),
FacetValue::String("sports".to_string()),
);
engine
.index_text(b"ns", b"doc1", "sports update today", facets2)
.unwrap();
let results = engine.search_text(b"ns", "breaking", 5).unwrap();
assert!(
results.is_empty(),
"old text 'breaking' should not match after re-index"
);
let results2 = engine.search_text(b"ns", "sports", 5).unwrap();
assert_eq!(results2.len(), 1);
assert_eq!(results2[0].doc_id, b"doc1");
let text_ns = edgestore::text_namespace(b"ns");
let raw = engine.get(&text_ns, b"doc1").unwrap().unwrap();
let record = edgestore::decode_text_record(&raw).unwrap();
assert_eq!(
record.facets.get("category"),
Some(&FacetValue::String("sports".to_string()))
);
}
#[test]
fn test_crash_recovery_rebuilds_stale_sidecar() {
let dir = TempDir::new().unwrap();
{
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine.flush().unwrap();
}
{
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc2",
"hello foo",
std::collections::HashMap::new(),
)
.unwrap();
}
{
let engine = open_engine(&dir);
let results = engine.search_text(b"ns", "hello", 5).unwrap();
assert_eq!(
results.len(),
2,
"crash recovery must rebuild stale sidecar so both docs are searchable"
);
assert!(results.iter().any(|r| r.doc_id == b"doc1"));
assert!(results.iter().any(|r| r.doc_id == b"doc2"));
}
}
#[test]
fn test_no_rebuild_when_sidecar_fresh() {
let dir = TempDir::new().unwrap();
{
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine.flush().unwrap();
engine
.index_text(
b"ns",
b"doc2",
"hello foo",
std::collections::HashMap::new(),
)
.unwrap();
engine.flush().unwrap();
}
{
let engine = open_engine(&dir);
let results = engine.search_text(b"ns", "hello", 5).unwrap();
assert_eq!(
results.len(),
2,
"fresh sidecar should contain both docs without rebuild"
);
}
}
#[test]
fn test_reindex_after_reload_removes_old_terms() {
let dir = TempDir::new().unwrap();
{
let mut engine = open_engine(&dir);
engine
.index_text(
b"ns",
b"doc1",
"hello world",
std::collections::HashMap::new(),
)
.unwrap();
engine.flush().unwrap();
}
{
let mut engine = open_engine(&dir);
engine
.index_text(b"ns", b"doc1", "foo bar", std::collections::HashMap::new())
.unwrap();
let results_hello = engine.search_text(b"ns", "hello", 5).unwrap();
assert!(
results_hello.is_empty(),
"old term 'hello' must not match after re-indexing post-reload"
);
let results_foo = engine.search_text(b"ns", "foo", 5).unwrap();
assert_eq!(results_foo.len(), 1);
assert_eq!(results_foo[0].doc_id, b"doc1");
}
}