use std::collections::BTreeMap;
use grust_core::prelude::*;
use crate::{LadybugConfig, LadybugGraphMode, LadybugGraphStore, LadybugPath};
fn props(entries: &[(&str, Value)]) -> Props {
entries
.iter()
.map(|(key, value)| ((*key).to_string(), value.clone()))
.collect::<BTreeMap<_, _>>()
}
fn sample_graph() -> Graph {
Graph::new(
vec![
Node::new(
"Person",
"person:ada",
props(&[("name", Value::from("Ada"))]),
),
Node::new(
"Talk",
"talk:grust",
props(&[("title", Value::from("Grust"))]),
),
],
vec![
Edge::new(
"Presented By",
"person:ada",
"talk:grust",
props(&[("year", Value::from(2026_i64))]),
)
.with_id("edge:presented"),
],
)
}
#[tokio::test]
async fn put_graph_reads_nodes_edges_and_traverses() -> Result<()> {
let store = LadybugGraphStore::new(LadybugConfig::untyped())?;
store.bootstrap().await?;
let graph = sample_graph();
assert_eq!(
store.put_graph(&graph).await?,
LoadReport { nodes: 2, edges: 1 }
);
let node = store
.get_node(&NodeId::from("person:ada"))
.await?
.expect("node should be readable");
assert_eq!(node.label, Label::from("Person"));
assert_eq!(node.props.get("name"), Some(&Value::from("Ada")));
let edges = store
.get_edges(EdgeQuery {
from: Some(NodeId::from("person:ada")),
to: None,
label: Some(Label::from("Presented By")),
})
.await?;
assert_eq!(edges.len(), 1);
assert_eq!(edges[0].to, NodeId::from("talk:grust"));
assert_eq!(edges[0].props.get("year"), Some(&Value::from(2026_i64)));
let traversed = store
.traverse(
Traversal::from_node("person:ada")
.out("Presented By")
.to("Talk"),
)
.await?;
assert_eq!(traversed.len(), 1);
assert_eq!(traversed[0].id, NodeId::from("talk:grust"));
Ok(())
}
#[test]
fn ladybug_graph_mode_helpers_map_to_dynamic_schema() {
assert_eq!(
LadybugConfig::untyped().graph_mode(),
LadybugGraphMode::Untyped
);
assert!(LadybugConfig::untyped().dynamic_schema);
assert_eq!(LadybugConfig::typed().graph_mode(), LadybugGraphMode::Typed);
assert!(!LadybugConfig::typed().dynamic_schema);
}
#[tokio::test]
async fn applies_schema_before_typed_graph_write() -> Result<()> {
let store = LadybugGraphStore::new(LadybugConfig::typed())?;
let schema = GraphSchema::builder()
.node("Person", vec![])
.node("Talk", vec![])
.edge(
"Presented By",
vec![Label::from("Person")],
vec![Label::from("Talk")],
vec![],
)
.build();
let report = store.put_typed_graph(&schema, &sample_graph()).await?;
assert_eq!(report, LoadReport { nodes: 2, edges: 1 });
Ok(())
}
#[tokio::test]
async fn typed_mode_requires_schema_before_writes() -> Result<()> {
let store = LadybugGraphStore::new(LadybugConfig::typed())?;
let err = store
.put_graph(&sample_graph())
.await
.expect_err("typed Ladybug mode should require schema first");
assert!(err.to_string().contains("requires apply_schema"));
Ok(())
}
#[tokio::test]
async fn applied_schema_validates_later_untyped_writes() -> Result<()> {
let store = LadybugGraphStore::in_memory()?;
let schema = GraphSchema::builder()
.node("Person", vec![Field::required("name", FieldType::String)])
.build();
store.apply_schema(&schema).await?;
let err = store
.put_node(&Node::new("Person", "person:ada", Props::new()))
.await
.expect_err("applied schema should validate later writes");
assert!(err.to_string().contains("missing required field 'name'"));
Ok(())
}
#[tokio::test]
async fn persists_to_directory() -> Result<()> {
let tempdir =
tempfile::tempdir().map_err(|err| GrustError::Backend(format!("tempdir error: {err}")))?;
let path = tempdir.path().join("ladybug");
{
let store = LadybugGraphStore::new(LadybugConfig {
path: LadybugPath::Directory(path.clone()),
..LadybugConfig::default()
})?;
store.put_graph(&sample_graph()).await?;
}
{
let store = LadybugGraphStore::open(path)?;
let node = store.get_node(&NodeId::from("talk:grust")).await?;
assert!(node.is_some());
}
Ok(())
}
#[tokio::test]
async fn clear_removes_managed_graph() -> Result<()> {
let store = LadybugGraphStore::in_memory()?;
store.put_graph(&sample_graph()).await?;
store.clear().await?;
assert!(store.get_node(&NodeId::from("person:ada")).await?.is_none());
Ok(())
}
#[tokio::test]
async fn clear_preserves_applied_schema_tables_in_typed_mode() -> Result<()> {
let store = LadybugGraphStore::new(LadybugConfig::typed())?;
let schema = GraphSchema::builder()
.node("Person", vec![])
.node("Talk", vec![])
.edge(
"Presented By",
vec![Label::from("Person")],
vec![Label::from("Talk")],
vec![],
)
.build();
store.put_typed_graph(&schema, &sample_graph()).await?;
store.clear().await?;
let report = store.put_graph(&sample_graph()).await?;
assert_eq!(report, LoadReport { nodes: 2, edges: 1 });
Ok(())
}
#[cfg(feature = "arrow")]
mod arrow_tests {
use std::sync::Arc;
use arrow::{
array::{Array as _, Int64Array, StringArray},
datatypes::{DataType, Field, Schema},
ipc::{reader::StreamReader, writer::StreamWriter},
record_batch::RecordBatch,
};
use super::*;
fn arrow_err(context: &str, err: impl std::fmt::Display) -> GrustError {
GrustError::Backend(format!("{context}: {err}"))
}
fn ipc_bytes(batch: &RecordBatch) -> Result<Vec<u8>> {
let mut data = Vec::new();
{
let cursor = std::io::Cursor::new(&mut data);
let mut writer = StreamWriter::try_new(cursor, batch.schema().as_ref())
.map_err(|err| arrow_err("Arrow IPC writer", err))?;
writer
.write(batch)
.map_err(|err| arrow_err("Arrow IPC write", err))?;
writer
.finish()
.map_err(|err| arrow_err("Arrow IPC finish", err))?;
}
Ok(data)
}
fn collect_string_column(chunks: &[Vec<u8>], column: usize) -> Result<Vec<String>> {
let mut values = Vec::new();
for chunk in chunks {
let reader = StreamReader::try_new(std::io::Cursor::new(chunk), None)
.map_err(|err| arrow_err("Arrow IPC reader", err))?;
for batch in reader {
let batch = batch.map_err(|err| arrow_err("Arrow batch", err))?;
let array = batch
.column(column)
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| GrustError::Schema("query column is not string".into()))?;
values.extend((0..array.len()).map(|index| array.value(index).to_string()));
}
}
Ok(values)
}
fn collect_i64_column(chunks: &[Vec<u8>], column: usize) -> Result<Vec<i64>> {
let mut values = Vec::new();
for chunk in chunks {
let reader = StreamReader::try_new(std::io::Cursor::new(chunk), None)
.map_err(|err| arrow_err("Arrow IPC reader", err))?;
for batch in reader {
let batch = batch.map_err(|err| arrow_err("Arrow batch", err))?;
let array = batch
.column(column)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| GrustError::Schema("query column is not i64".into()))?;
values.extend((0..array.len()).map(|index| array.value(index)));
}
}
Ok(values)
}
#[test]
fn arrow_ipc_node_table_queries_through_ladybug() -> Result<()> {
let store = LadybugGraphStore::in_memory()?;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec!["Ada", "Grace"])),
],
)
.map_err(|err| arrow_err("Arrow node batch", err))?;
store.register_arrow_ipc_node_table("Person", &ipc_bytes(&batch)?)?;
let chunks =
store.query_arrow_ipc("MATCH (p:Person) RETURN p.name ORDER BY p.id;", 1024)?;
assert_eq!(collect_string_column(&chunks, 0)?, vec!["Ada", "Grace"]);
Ok(())
}
#[test]
fn arrow_ipc_relationship_table_queries_through_ladybug() -> Result<()> {
let store = LadybugGraphStore::in_memory()?;
let node_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
let nodes = RecordBatch::try_new(node_schema, vec![Arc::new(Int64Array::from(vec![0, 1]))])
.map_err(|err| arrow_err("Arrow node batch", err))?;
store.register_arrow_ipc_node_table("Person", &ipc_bytes(&nodes)?)?;
let rel_schema = Arc::new(Schema::new(vec![
Field::new("from", DataType::Int64, false),
Field::new("to", DataType::Int64, false),
Field::new("weight", DataType::Int64, false),
]));
let rels = RecordBatch::try_new(
rel_schema,
vec![
Arc::new(Int64Array::from(vec![0, 1])),
Arc::new(Int64Array::from(vec![1, 0])),
Arc::new(Int64Array::from(vec![7, 9])),
],
)
.map_err(|err| arrow_err("Arrow relationship batch", err))?;
store.register_arrow_ipc_rel_table("Knows", &ipc_bytes(&rels)?, "Person", "Person")?;
let chunks = store.query_arrow_ipc(
"MATCH (a:Person)-[r:Knows]->(b:Person) \
RETURN r.weight ORDER BY a.id, b.id;",
1024,
)?;
assert_eq!(collect_i64_column(&chunks, 0)?, vec![7, 9]);
Ok(())
}
}