use std::io;
use std::num::{NonZeroU32, NonZeroUsize};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, TryRecvError};
use std::time::Duration;
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::{Value, json, value::RawValue};
use super::{ToolResult, dispatch};
use crate::mcp::{rpc, stdio, types};
use crate::serve::{frame, protocol, service::Service, storage::Storage};
type TestResult = Result<(), Box<dyn std::error::Error>>;
const CAP: NonZeroU32 = NonZeroU32::MIN.saturating_add(1_048_575);
const ROWS: NonZeroU32 = NonZeroU32::MIN.saturating_add(63);
pub(super) struct Fixture {
pub(super) sender: stdio::NativeSender,
service: Arc<Service>,
socket: PathBuf,
stop: mpsc::Sender<()>,
task: std::thread::JoinHandle<io::Result<()>>,
calls: Arc<AtomicUsize>,
directory: tempfile::TempDir,
}
impl Fixture {
pub(super) fn start() -> Result<Self, Box<dyn std::error::Error>> {
let directory = tempfile::tempdir()?;
let 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::KeepAll,
},
NodeCacheBudget::Unlimited,
fixture_long_wait_interval(),
)?;
let service = Arc::new(Service::new(storage, ROWS)?);
let socket = directory.path().join("service.sock");
let listener = UnixListener::bind(&socket)?;
let (stop, stopped) = mpsc::channel();
let calls = Arc::new(AtomicUsize::new(0));
let served = Arc::clone(&service);
let requested = Arc::clone(&calls);
let task = std::thread::spawn(move || {
loop {
let (mut stream, _) = listener.accept()?;
match stopped.try_recv() {
Ok(()) | Err(TryRecvError::Disconnected) => return Ok(()),
Err(TryRecvError::Empty) => {}
}
while let Some(bytes) = frame::read(&mut stream, CAP)? {
let request: protocol::Request =
serde_json::from_slice(&bytes).map_err(io::Error::other)?;
let response = if matches!(&request.operation, protocol::Operation::Hello(_)) {
json!({"id":request.id,"result":{
"protocol":protocol::PROTOCOL,
"capabilities":[protocol::RECEIPT_CAPABILITY],
"max_frame":CAP,"max_page":ROWS,
}})
} else {
requested.fetch_add(1, Ordering::SeqCst);
let budget = NonZeroUsize::new(1_048_576)
.ok_or_else(|| io::Error::other("fixture budget is zero"))?;
match served.dispatch_text(request.operation, budget) {
Ok(raw) => json!({"id":request.id,"result":raw}),
Err(refusal) => json!({"id":request.id,"error":refusal}),
}
};
frame::write(
&mut stream,
&serde_json::to_vec(&response).map_err(io::Error::other)?,
CAP,
)?;
}
}
});
let sender = stdio::NativeSender::new(socket.clone(), Duration::from_secs(1));
Ok(Self {
sender,
service,
socket,
stop,
task,
calls,
directory,
})
}
pub(super) fn current(&self) -> Result<protocol::Version, Box<dyn std::error::Error>> {
let answer =
self.service
.dispatch(protocol::Operation::StoreStats(protocol::StoreRead {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
}))?;
Ok(serde_json::from_value(
answer["selection"]["version"].clone(),
)?)
}
pub(super) fn requests(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
pub(super) fn mutate(
&self,
branch: &str,
request_id: &str,
change: protocol::Change,
) -> Result<protocol::Version, Box<dyn std::error::Error>> {
let capture =
self.service
.dispatch(protocol::Operation::StoreStats(protocol::StoreRead {
selection: protocol::Selection {
branch: branch.to_owned(),
version: None,
},
}))?;
let answer = self
.service
.dispatch(protocol::Operation::Mutate(protocol::Mutation {
request_id: request_id.to_owned(),
branch: branch.to_owned(),
expected: serde_json::from_value(capture["selection"]["version"].clone())?,
change,
}))?;
Ok(serde_json::from_value(answer["version"].clone())?)
}
pub(super) fn fork(&self, name: &str) -> TestResult {
self.mutate(
"main",
&format!("fork-{name}"),
protocol::Change::Fork(protocol::Fork {
name: name.to_owned(),
from: self.current()?,
}),
)?;
Ok(())
}
fn call(
&mut self,
name: &str,
mut fields: Value,
) -> Result<ToolResult, Box<dyn std::error::Error>> {
let expected = self.current()?;
let object = fields
.as_object_mut()
.ok_or_else(|| io::Error::other("input must be an object"))?;
object.insert(
"request_id".to_owned(),
Value::String(format!("request-{name}")),
);
object.insert("branch".to_owned(), Value::String("main".to_owned()));
object.insert("expected".to_owned(), serde_json::to_value(expected)?);
let call = rpc::ToolCall {
metadata: None,
name: name.to_owned(),
arguments: serde_json::value::to_raw_value(&fields)?,
};
let before = self.calls.load(Ordering::SeqCst);
let result = dispatch(&mut self.sender, &call)?;
assert_eq!(self.calls.load(Ordering::SeqCst), before + 1);
Ok(result)
}
pub(super) fn seed(&self) -> TestResult {
for change in [
protocol::Change::Collection(protocol::CollectionWrite {
collection: "records".to_owned(),
schema: None,
}),
protocol::Change::PutMany(protocol::DocumentImport {
collection: "records".to_owned(),
documents: vec![
protocol::ImportedDocument {
document_id: "a".to_owned(),
document: RawValue::from_string(r#"{"n":1}"#.to_owned())?,
},
protocol::ImportedDocument {
document_id: "b".to_owned(),
document: RawValue::from_string(r#"{"n":2}"#.to_owned())?,
},
],
}),
protocol::Change::RelationPut(protocol::Relation {
relation: "links".to_owned(),
from_collection: "records".to_owned(),
to_collection: "records".to_owned(),
cardinality: protocol::Cardinality::ManyToMany,
}),
] {
self.service
.dispatch(protocol::Operation::Mutate(protocol::Mutation {
request_id: format!("seed-{}", self.current()?.sequence),
branch: "main".to_owned(),
expected: self.current()?,
change,
}))?;
}
Ok(())
}
pub(super) fn finish(self) -> TestResult {
drop(self.sender);
self.stop.send(())?;
drop(UnixStream::connect(&self.socket)?);
self.task
.join()
.map_err(|_| io::Error::other("fixture service task panicked"))??;
drop(self.service);
self.directory.close()?;
Ok(())
}
}
fn successful(result: &ToolResult, name: &str) -> TestResult {
assert!(!result.is_error, "{}", result.structured_content.get());
let schema = crate::mcp::schema::tools()
.into_iter()
.find(|tool| tool.name == name)
.ok_or_else(|| io::Error::other("tool schema missing"))?;
let validator = jsonschema::validator_for(schema.output_schema.as_value())?;
let output: Value = serde_json::from_str(result.structured_content.get())?;
validator
.validate(&output)
.map_err(|error| io::Error::other(error.to_string()))?;
Ok(())
}
#[test]
fn define_collection_uses_one_served_mutation_and_typed_output() -> TestResult {
let mut fixture = Fixture::start()?;
let result = fixture.call(
"define_collection",
json!({"collection":"records","schema":null}),
)?;
successful(&result, "define_collection")?;
let answer =
fixture
.service
.dispatch(protocol::Operation::Collections(protocol::CollectionPage {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
after: None,
limit: ROWS,
}))?;
assert_eq!(answer["entries"][0]["name"], "records");
fixture.finish()
}
#[test]
fn put_records_is_atomic_and_keeps_raw_number_spelling() -> TestResult {
let mut fixture = Fixture::start()?;
fixture.seed()?;
let raw: Box<RawValue> =
serde_json::from_str(r#"{"n":1e2,"identity":18446744073709551616001}"#)?;
let expected = serde_json::to_string(&fixture.current()?)?;
let arguments = format!(
r#"{{"request_id":"request-put_records","branch":"main","expected":{expected},"collection":"records","records":[{{"document_id":"large","document":{}}},{{"document_id":"second","document":{{"n":3}}}}]}}"#,
raw.get(),
);
let call = rpc::ToolCall {
metadata: None,
name: "put_records".to_owned(),
arguments: RawValue::from_string(arguments)?,
};
let before = fixture.requests();
let result = dispatch(&mut fixture.sender, &call)?;
assert_eq!(fixture.requests(), before + 1);
successful(&result, "put_records")?;
let answer = fixture
.service
.dispatch(protocol::Operation::Get(protocol::DocumentRead {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
collection: "records".to_owned(),
document_id: "large".to_owned(),
}))?;
assert_eq!(
answer["document"]["raw_json"],
r#"{"n":1e2,"identity":18446744073709551616001}"#
);
let second = fixture
.service
.dispatch(protocol::Operation::Get(protocol::DocumentRead {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
collection: "records".to_owned(),
document_id: "second".to_owned(),
}))?;
assert_eq!(second["document"]["raw_json"], r#"{"n":3}"#);
fixture.finish()
}
#[test]
fn delete_record_uses_served_delete_and_returns_the_receipt_version() -> TestResult {
let mut fixture = Fixture::start()?;
fixture.seed()?;
let result = fixture.call(
"delete_record",
json!({"record":{"collection":"records","document_id":"a"}}),
)?;
successful(&result, "delete_record")?;
let output: types::DeleteRecordOutput = serde_json::from_str(result.structured_content.get())?;
assert_eq!(output.version.0, fixture.current()?);
let answer = fixture
.service
.dispatch(protocol::Operation::Get(protocol::DocumentRead {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
collection: "records".to_owned(),
document_id: "a".to_owned(),
}))?;
assert!(answer["document"].is_null());
fixture.finish()
}
fn endpoints() -> Value {
json!({"relation":"links","from":{"collection":"records","document_id":"a"},"to":{"collection":"records","document_id":"b"}})
}
#[test]
fn link_records_passes_both_endpoints_and_properties_to_the_service() -> TestResult {
let mut fixture = Fixture::start()?;
fixture.seed()?;
let mut fields = endpoints();
fields["properties"] = json!({"label":"edge"});
let result = fixture.call("link_records", fields)?;
successful(&result, "link_records")?;
let answer = fixture
.service
.dispatch(protocol::Operation::EdgesOut(protocol::EdgePage {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
node: protocol::Node {
collection: "records".to_owned(),
document_id: "a".to_owned(),
},
relation: None,
after: None,
limit: ROWS,
}))?;
assert_eq!(answer["edges"][0]["properties"]["label"], "edge");
assert_eq!(answer["edges"][0]["from"]["document_id"], "a");
assert_eq!(answer["edges"][0]["to"]["document_id"], "b");
fixture.finish()
}
#[test]
fn unlink_records_removes_only_the_named_served_edge() -> TestResult {
let mut fixture = Fixture::start()?;
fixture.seed()?;
for destination in ["a", "b"] {
fixture
.service
.dispatch(protocol::Operation::Mutate(protocol::Mutation {
request_id: format!("seed-edge-{destination}"),
branch: "main".to_owned(),
expected: fixture.current()?,
change: protocol::Change::EdgePut(protocol::Edge {
relation: "links".to_owned(),
from: protocol::Node {
collection: "records".to_owned(),
document_id: "a".to_owned(),
},
to: protocol::Node {
collection: "records".to_owned(),
document_id: destination.to_owned(),
},
properties: json!({}),
}),
}))?;
}
let result = fixture.call("unlink_records", endpoints())?;
successful(&result, "unlink_records")?;
let answer = fixture
.service
.dispatch(protocol::Operation::EdgesOut(protocol::EdgePage {
selection: protocol::Selection {
branch: "main".to_owned(),
version: None,
},
node: protocol::Node {
collection: "records".to_owned(),
document_id: "a".to_owned(),
},
relation: None,
after: None,
limit: ROWS,
}))?;
assert_eq!(answer["edges"].as_array().map(Vec::len), Some(1));
assert_eq!(answer["edges"][0]["to"]["document_id"], "a");
fixture.finish()
}
#[test]
fn create_branch_uses_the_selected_from_version_and_reports_destination() -> TestResult {
let mut fixture = Fixture::start()?;
fixture.seed()?;
let from = fixture.current()?;
let result = fixture.call("create_branch", json!({"name":"feature","from":from}))?;
successful(&result, "create_branch")?;
let output: types::CreateBranchOutput = serde_json::from_str(result.structured_content.get())?;
assert_eq!(output.branch, "feature");
assert_eq!(output.request_id, "request-create_branch");
let answer =
fixture
.service
.dispatch(protocol::Operation::Collections(protocol::CollectionPage {
selection: protocol::Selection {
branch: "feature".to_owned(),
version: Some(output.version.0),
},
after: None,
limit: ROWS,
}))?;
assert_eq!(answer["entries"][0]["name"], "records");
fixture.finish()
}
const fn fixture_long_wait_interval() -> std::time::Duration {
std::time::Duration::from_secs(1)
}