use serde_json::json;
use super::{Arc, NonZeroU32, Service, TestResult, TestSource, response_limit, selection};
use crate::serve::protocol::{
Change, Metric, Mutation, Operation, Selection, StoreRead, VectorPage, VectorRead,
VectorSearch, VectorWrite, Version,
};
use crate::serve::storage::Storage;
fn with_vectors(test: impl FnOnce(Storage) -> TestResult) -> TestResult {
super::with_list_storage(|mut storage, _| {
for (id, value) in [("a", 4.0), ("b", 5.0), ("c", 3.0), ("d", 2.0), ("e", 1.0)] {
let expected = storage.capture(&selection())?.version;
storage.mutate(&Mutation {
request_id: format!("vector-{id}"),
branch: "main".to_owned(),
expected,
change: Change::VectorPut(VectorWrite {
collection: "records".to_owned(),
document_id: id.to_owned(),
field: "embedding".to_owned(),
vector: vec![value],
}),
})?;
}
test(storage)
})
}
fn get(id: &str) -> Operation {
Operation::VectorGet(VectorRead {
selection: selection(),
collection: "records".to_owned(),
document_id: id.to_owned(),
field: "embedding".to_owned(),
})
}
fn page(
after: Option<&str>,
version: Option<Version>,
rows: u32,
) -> Result<Operation, Box<dyn std::error::Error>> {
Ok(Operation::VectorList(VectorPage {
selection: Selection {
branch: "main".to_owned(),
version,
},
collection: "records".to_owned(),
field: "embedding".to_owned(),
after: after.map(str::to_owned),
limit: NonZeroU32::new(rows).ok_or("zero vector page limit")?,
}))
}
fn search() -> Result<Operation, Box<dyn std::error::Error>> {
Ok(Operation::VectorSearch(VectorSearch {
selection: selection(),
collection: "records".to_owned(),
field: "embedding".to_owned(),
query: vec![1.0],
k: NonZeroU32::new(2).ok_or("zero vector search limit")?,
metric: Metric::Dot,
}))
}
fn fields() -> Operation {
Operation::VectorFields(StoreRead {
selection: selection(),
})
}
fn operations() -> Result<[Operation; 4], Box<dyn std::error::Error>> {
Ok([get("b"), page(None, None, 2)?, search()?, fields()])
}
#[test]
fn denied_vector_is_byte_equal_to_an_absent_vector() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::denying(&[("records", "b")]));
let service = Service::with_access_source(storage, NonZeroU32::MIN, None, source)?;
let denied = service.dispatch_bounded(get("b"), response_limit()?)?;
let absent = service.dispatch_bounded(get("absent"), response_limit()?)?;
assert!(denied["vector"].is_null());
assert_eq!(serde_json::to_vec(&denied)?, serde_json::to_vec(&absent)?);
Ok(())
})
}
#[test]
fn vector_pages_continue_from_visible_rows_without_repeating_or_skipping() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::denying(&[("records", "b"), ("records", "d")]));
let service = Service::with_access_source(
storage,
NonZeroU32::new(2).ok_or("zero page limit")?,
None,
source,
)?;
let first = service.dispatch_bounded(page(None, None, 2)?, response_limit()?)?;
assert_eq!(
first["vectors"],
json!([
{"document_id":"a", "vector":[4.0]},
{"document_id":"c", "vector":[3.0]},
])
);
assert_eq!(first["next_after"], "c");
let at: Version = serde_json::from_value(first["selection"]["version"].clone())?;
let second = service.dispatch_bounded(page(Some("c"), Some(at), 2)?, response_limit()?)?;
assert_eq!(
second["vectors"],
json!([{"document_id":"e", "vector":[1.0]}])
);
assert!(second["next_after"].is_null());
assert_eq!(second["selection"], first["selection"]);
let owner = service.dispatch_bounded(get("a"), response_limit()?)?;
assert_eq!(owner["vector"], json!([4.0]));
Ok(())
})
}
#[test]
fn vector_page_cursor_never_names_a_denied_raw_boundary() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::denying(&[("records", "b")]));
let service = Service::with_access_source(storage, NonZeroU32::MIN, None, source)?;
let first = service.dispatch_bounded(page(None, None, 1)?, response_limit()?)?;
assert_eq!(
first["vectors"],
json!([{"document_id":"a", "vector":[4.0]}])
);
assert_eq!(first["next_after"], "a");
let at: Version = serde_json::from_value(first["selection"]["version"].clone())?;
let second = service.dispatch_bounded(page(Some("a"), Some(at), 1)?, response_limit()?)?;
assert_eq!(
second["vectors"],
json!([{"document_id":"c", "vector":[3.0]}])
);
assert_eq!(second["next_after"], "c");
Ok(())
})
}
#[test]
fn denied_vector_run_over_the_page_scan_cap_is_refused_by_name() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::denying(&[
("records", "a"),
("records", "b"),
("records", "c"),
("records", "d"),
("records", "e"),
]));
let service = Service::with_access_source(storage, NonZeroU32::MIN, None, source)?;
for bounded in [false, true] {
let result = if bounded {
service.dispatch_bounded(page(None, None, 1)?, response_limit()?)
} else {
service.dispatch(page(None, None, 1)?)
};
assert_eq!(
result
.err()
.ok_or("over-cap denied run returned vectors")?
.code,
"access_scan_limit"
);
}
Ok(())
})
}
#[test]
fn vector_search_replaces_denied_best_hit_and_counts_only_visible_vectors() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::denying(&[("records", "b"), ("records", "d")]));
let service = Service::with_access_source(
storage,
NonZeroU32::new(2).ok_or("zero page limit")?,
None,
source,
)?;
let answer = service.dispatch_bounded(search()?, response_limit()?)?;
assert_eq!(
answer["hits"],
json!([
{"document_id":"a", "score":4.0}, {"document_id":"c", "score":3.0},
])
);
assert_eq!(answer["scanned"], 3);
Ok(())
})
}
#[test]
fn vector_field_counts_only_visible_vectors_and_keeps_the_dimension() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::denying(&[("records", "b"), ("records", "d")]));
let service = Service::with_access_source(storage, NonZeroU32::MIN, None, source)?;
let answer = service.dispatch_bounded(fields(), response_limit()?)?;
assert_eq!(
answer["fields"],
json!([
{"collection":"records", "field":"embedding", "dimensions":1, "count":3},
])
);
Ok(())
})
}
#[test]
fn vector_reads_resolve_access_once_and_propagate_source_failure() -> TestResult {
with_vectors(|storage| {
let source = Arc::new(TestSource::unavailable());
let service = Service::with_access_source(
storage,
NonZeroU32::new(2).ok_or("zero page limit")?,
None,
source.clone(),
)?;
for (index, operation) in operations()?.into_iter().enumerate() {
let failure = service
.dispatch_bounded(operation, response_limit()?)
.err()
.ok_or("unavailable source served vectors")?;
assert_eq!(failure.code, "access_unavailable");
assert_eq!(source.calls(), index + 1);
}
for (index, operation) in operations()?.into_iter().enumerate() {
let failure = service
.dispatch(operation)
.err()
.ok_or("unavailable source served unbounded vectors")?;
assert_eq!(failure.code, "access_unavailable");
assert_eq!(source.calls(), index + 5);
}
Ok(())
})
}