use haematite::SyncCounts;
use super::batch_test_support::{ONE_WRITE, TestResult, apply, at, record_id, store, vector};
use super::{Failure, Storage, reads, vectors};
use crate::serve::protocol::{Change, Mutation, ReceiptQuery, VectorImport, VectorWrite};
fn many(vectors: Vec<VectorWrite>) -> Change {
Change::VectorPutMany(VectorImport { vectors })
}
fn refused_whole(
storage: &mut Storage,
vectors: Vec<VectorWrite>,
) -> Result<Failure, Box<dyn std::error::Error>> {
let before = at(storage)?;
let counts = storage.sync_counts();
let failure = apply(storage, "refused", many(vectors))
.err()
.ok_or("invalid vector batch accepted")?;
assert_eq!(at(storage)?.version, before.version);
assert_eq!(storage.sync_counts().since(counts), SyncCounts::default());
assert_eq!(
reads::test_support::count_keys(&storage.reader, &at(storage)?, b"v\0", 3)?,
1,
);
assert!(
storage
.receipt(&ReceiptQuery {
branch: "main".to_owned(),
request_id: "refused".to_owned(),
})?
.is_none()
);
Ok(failure)
}
#[test]
fn five_hundred_vectors_share_one_commit_and_the_control_flushes() -> TestResult {
let (_directory, mut storage) = store(500)?;
let before = at(&storage)?;
let counts = storage.sync_counts();
let receipt = apply(
&mut storage,
"vectors",
many((0..500).map(vector).collect()),
)?;
assert_eq!(receipt.version.sequence, before.version.sequence + 1);
assert_eq!(receipt.version, at(&storage)?.version);
assert_eq!(storage.sync_counts().since(counts), ONE_WRITE);
let access = crate::serve::access::request(&crate::serve::access::Standalone)?;
for index in 0..500 {
assert_eq!(
vectors::get(
&access,
&storage.reader,
&at(&storage)?,
"records",
"supplied",
&record_id(index),
)?,
Some(vec![1.0, 2.0, 3.0]),
);
}
let audit = serde_json::to_value(Storage::audit_counts(
&access,
&storage.reader,
&at(&storage)?,
)?)?;
assert!(
audit["buckets"]
.as_array()
.ok_or("no count buckets")?
.iter()
.all(|row| row["result"] == "agree")
);
Ok(())
}
#[test]
fn dimension_failure_at_137_refuses_all_five_hundred_without_flushes() -> TestResult {
let (_directory, mut storage) = store(500)?;
let mut batch: Vec<_> = (0..500).map(vector).collect();
batch[137].vector.truncate(2);
let failure = refused_whole(&mut storage, batch)?;
assert_eq!(failure.code, "vector_dimension");
assert_eq!(failure.path.as_deref(), Some("/vectors/137/vector"));
Ok(())
}
#[test]
fn duplicate_record_and_field_refuses_whole() -> TestResult {
let (_directory, mut storage) = store(2)?;
let failure = refused_whole(&mut storage, vec![vector(0), vector(1), vector(0)])?;
assert_eq!(failure.code, "import_duplicate");
assert_eq!(failure.path.as_deref(), Some("/vectors/2"));
Ok(())
}
#[test]
fn undeclared_field_refuses_whole() -> TestResult {
let (_directory, mut storage) = store(2)?;
let mut second = vector(1);
second.field = "unknown".to_owned();
let failure = refused_whole(&mut storage, vec![vector(0), second])?;
assert_eq!(failure.code, "vector_dimension_missing");
assert_eq!(failure.path.as_deref(), Some("/vectors/1/field"));
Ok(())
}
#[test]
fn missing_record_refuses_whole() -> TestResult {
let (_directory, mut storage) = store(2)?;
let failure = refused_whole(&mut storage, vec![vector(0), vector(2)])?;
assert_eq!(failure.code, "document_missing");
assert_eq!(failure.path.as_deref(), Some("/vectors/1/document_id"));
Ok(())
}
#[test]
fn empty_and_non_finite_vectors_refuse_whole() -> TestResult {
for (members, code) in [
(vec![], "vector_empty"),
(vec![f32::INFINITY], "vector_non_finite"),
] {
let (_directory, mut storage) = store(2)?;
let mut second = vector(1);
second.vector = members;
let failure = refused_whole(&mut storage, vec![vector(0), second])?;
assert_eq!(failure.code, code);
assert_eq!(failure.path.as_deref(), Some("/vectors/1/vector"));
}
Ok(())
}
#[test]
fn empty_and_over_bound_batches_reuse_document_import_refusals() -> TestResult {
let (_directory, mut storage) = store(2)?;
let failure = refused_whole(&mut storage, vec![])?;
assert_eq!(failure.code, "import_empty");
assert_eq!(failure.path.as_deref(), Some("/vectors"));
let failure = refused_whole(&mut storage, (0..513).map(vector).collect())?;
assert_eq!(failure.code, "import_limit");
assert_eq!(failure.path.as_deref(), Some("/vectors/512"));
let request = Mutation {
request_id: "unbounded".to_owned(),
branch: "main".to_owned(),
expected: at(&storage)?.version,
change: many(vec![vector(0)]),
};
let counts = storage.sync_counts();
let failure = storage
.mutate(&request)
.err()
.ok_or("unbounded import accepted")?;
assert_eq!(failure.code, "import_limit_required");
assert_eq!(storage.sync_counts().since(counts), SyncCounts::default());
Ok(())
}
#[test]
fn replay_reuses_the_batch_receipt_without_another_commit() -> TestResult {
let (_directory, mut storage) = store(2)?;
let request = Mutation {
request_id: "vectors".to_owned(),
branch: "main".to_owned(),
expected: at(&storage)?.version,
change: many(vec![vector(1), vector(0)]),
};
let bound = std::num::NonZeroU32::MIN.saturating_add(7);
let first = storage.mutate_bounded(&request, bound)?;
let counts = storage.sync_counts();
let second = storage.mutate_bounded(&request, bound)?;
assert_eq!(first.version, second.version);
assert_eq!(first.fingerprint, second.fingerprint);
assert_eq!(storage.sync_counts().since(counts).record_installs, 0);
Ok(())
}
#[test]
fn embedded_field_refuses_supplied_batch_without_flushes() -> TestResult {
use std::num::NonZeroUsize;
use super::embedded::{Declaration, ModelName, VectorKind};
use crate::serve::protocol::{Metric, VectorDeclaration};
let (_directory, mut storage) = store(2)?;
apply(
&mut storage,
"embedded",
Change::VectorDeclare(VectorDeclaration {
collection: "records".to_owned(),
field: "embedded".to_owned(),
declaration: Declaration {
kind: VectorKind::Dense,
source: "text".to_owned(),
model: ModelName {
name: "model".to_owned(),
revision: "revision".to_owned(),
manifest_sha256: "a".repeat(64),
},
dimensions: NonZeroUsize::MIN.saturating_add(2),
distance: Metric::Cosine,
},
}),
)?;
let mut second = vector(1);
second.field = "embedded".to_owned();
let failure = refused_whole(&mut storage, vec![vector(0), second])?;
assert_eq!(failure.code, "vector_field_embedded");
assert_eq!(failure.path.as_deref(), Some("/vectors/1/field"));
Ok(())
}
#[test]
fn replacing_vectors_keeps_counts_and_publishes_once() -> TestResult {
let (_directory, mut storage) = store(2)?;
apply(&mut storage, "first", many(vec![vector(1), vector(0)]))?;
let counts = storage.sync_counts();
let mut replacement = vector(1);
replacement.vector = vec![4.0, 5.0, 6.0];
apply(&mut storage, "second", many(vec![replacement, vector(0)]))?;
assert_eq!(storage.sync_counts().since(counts), ONE_WRITE);
let access = crate::serve::access::request(&crate::serve::access::Standalone)?;
assert_eq!(
vectors::get(
&access,
&storage.reader,
&at(&storage)?,
"records",
"supplied",
&record_id(1),
)?,
Some(vec![4.0, 5.0, 6.0]),
);
assert_eq!(
reads::test_support::count_keys(&storage.reader, &at(&storage)?, b"v\0", 3)?,
3,
);
let audit = serde_json::to_value(Storage::audit_counts(
&access,
&storage.reader,
&at(&storage)?,
)?)?;
assert!(
audit["buckets"]
.as_array()
.ok_or("no count buckets")?
.iter()
.all(|row| row["result"] == "agree")
);
Ok(())
}
#[test]
fn batch_vectors_and_receipt_survive_reopen_from_scalar_only_data() -> TestResult {
let (directory, storage) = store(2)?;
let previous = at(&storage)?.version;
drop(storage);
let mut storage = Storage::open(
&directory.path().join("database"),
haematite::store::NodeCacheBudget::Unlimited,
std::time::Duration::from_secs(1),
)?;
assert_eq!(at(&storage)?.version, previous);
let receipt = apply(&mut storage, "vectors", many(vec![vector(1), vector(0)]))?;
drop(storage);
let mut storage = Storage::open(
&directory.path().join("database"),
haematite::store::NodeCacheBudget::Unlimited,
std::time::Duration::from_secs(1),
)?;
assert_eq!(at(&storage)?.version, receipt.version);
let access = crate::serve::access::request(&crate::serve::access::Standalone)?;
for index in 0..2 {
assert_eq!(
vectors::get(
&access,
&storage.reader,
&at(&storage)?,
"records",
"supplied",
&record_id(index),
)?,
Some(vec![1.0, 2.0, 3.0]),
);
}
let replay = storage
.receipt(&ReceiptQuery {
branch: "main".to_owned(),
request_id: "vectors".to_owned(),
})?
.ok_or("batch receipt missing after reopen")?;
assert_eq!(replay.version, receipt.version);
assert_eq!(replay.fingerprint, receipt.fingerprint);
Ok(())
}