use std::error::Error;
use std::num::NonZeroU32;
use haematite::SyncCounts;
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::value::RawValue;
use super::{Captured, Failure, Receipt, Storage};
use crate::serve::protocol::{
Change, CollectionWrite, DocumentImport, ImportedDocument, Mutation, Selection, VectorWrite,
};
pub(super) type TestResult = Result<(), Box<dyn Error>>;
pub(super) const ONE_WRITE: SyncCounts = SyncCounts {
record_installs: 1,
barrier_device_flushes: 1,
ref_file_full_flushes: 1,
directory_full_flushes: 1,
};
pub(super) fn at(storage: &Storage) -> Result<Captured, Failure> {
storage.capture(&Selection {
branch: "main".to_owned(),
version: None,
})
}
pub(super) fn apply(storage: &mut Storage, id: &str, change: Change) -> Result<Receipt, Failure> {
storage.mutate_bounded(
&Mutation {
request_id: id.to_owned(),
branch: "main".to_owned(),
expected: at(storage)?.version,
change,
},
NonZeroU32::MIN.saturating_add(7),
)
}
pub(super) fn record_id(index: usize) -> String {
format!("record-{index:03}")
}
pub(super) fn vector(index: usize) -> VectorWrite {
VectorWrite {
collection: "records".to_owned(),
document_id: record_id(index),
field: "supplied".to_owned(),
vector: vec![1.0, 2.0, 3.0],
}
}
pub(super) fn store(rows: usize) -> Result<(tempfile::TempDir, Storage), Box<dyn Error>> {
let directory = tempfile::tempdir()?;
let mut storage = Storage::create(
DatabaseConfig {
data_dir: directory.path().join("database"),
shard_count: 1,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::reclaim(0),
},
NodeCacheBudget::Unlimited,
std::time::Duration::from_secs(1),
)?;
apply(
&mut storage,
"collection",
Change::Collection(CollectionWrite {
collection: "records".to_owned(),
schema: None,
}),
)?;
let mut documents = (0..rows)
.map(|index| {
Ok(ImportedDocument {
document_id: record_id(index),
document: RawValue::from_string("{}".to_owned())?,
})
})
.collect::<Result<Vec<_>, serde_json::Error>>()?;
documents.push(ImportedDocument {
document_id: "anchor".to_owned(),
document: RawValue::from_string("{}".to_owned())?,
});
apply(
&mut storage,
"records",
Change::PutMany(DocumentImport {
collection: "records".to_owned(),
documents,
}),
)?;
let mut declaration = vector(0);
declaration.document_id = "anchor".to_owned();
apply(&mut storage, "supplied", Change::VectorPut(declaration))?;
Ok((directory, storage))
}