use reifydb_core::{
common::CommitVersion,
encoded::shape::{RowShape, RowShapeField},
interface::{
catalog::{flow::FlowNodeId, id::TableId, shape::ShapeId},
change::{Change, ChangeOrigin, Diff},
},
row::Row,
value::column::columns::Columns,
};
use reifydb_type::value::{Value, datetime::DateTime, row_number::RowNumber, r#type::Type};
pub struct TestRowBuilder {
row_number: RowNumber,
values: Vec<Value>,
shape: Option<RowShape>,
}
impl TestRowBuilder {
pub fn new(row_number: impl Into<RowNumber>) -> Self {
Self {
row_number: row_number.into(),
values: Vec::new(),
shape: None,
}
}
pub fn with_values(mut self, values: Vec<Value>) -> Self {
self.values = values;
self
}
pub fn add_value(mut self, value: Value) -> Self {
self.values.push(value);
self
}
pub fn with_shape(mut self, shape: RowShape) -> Self {
self.shape = Some(shape);
self
}
pub fn build(self) -> Row {
let shape = if let Some(shape) = self.shape {
shape
} else {
let fields: Vec<RowShapeField> = self
.values
.iter()
.enumerate()
.map(|(i, v)| RowShapeField::unconstrained(format!("field{}", i), v.get_type()))
.collect();
RowShape::new(fields)
};
let mut encoded = shape.allocate();
shape.set_values(&mut encoded, &self.values);
Row {
number: self.row_number,
encoded,
shape,
}
}
}
pub struct TestChangeBuilder {
origin: ChangeOrigin,
diffs: Vec<Diff>,
version: CommitVersion,
changed_at: DateTime,
}
impl Default for TestChangeBuilder {
fn default() -> Self {
Self::new()
}
}
impl TestChangeBuilder {
pub fn new() -> Self {
Self {
origin: ChangeOrigin::Shape(ShapeId::Table(TableId(1))),
diffs: Vec::new(),
version: CommitVersion(1),
changed_at: DateTime::default(),
}
}
pub fn changed_by_shape(mut self, shape: ShapeId) -> Self {
self.origin = ChangeOrigin::Shape(shape);
self
}
pub fn changed_by_node(mut self, node: FlowNodeId) -> Self {
self.origin = ChangeOrigin::Flow(node);
self
}
pub fn with_version(mut self, version: CommitVersion) -> Self {
self.version = version;
self
}
pub fn with_changed_at(mut self, changed_at: DateTime) -> Self {
self.changed_at = changed_at;
self
}
pub fn insert(mut self, row: Row) -> Self {
self.diffs.push(Diff::Insert {
post: Columns::from_row(&row),
});
self
}
pub fn insert_row(self, row_number: impl Into<RowNumber>, values: Vec<Value>) -> Self {
let row = TestRowBuilder::new(row_number).with_values(values).build();
self.insert(row)
}
pub fn update(mut self, pre: Row, post: Row) -> Self {
self.diffs.push(Diff::Update {
pre: Columns::from_row(&pre),
post: Columns::from_row(&post),
});
self
}
pub fn update_row(
self,
row_number: impl Into<RowNumber>,
pre_values: Vec<Value>,
post_values: Vec<Value>,
) -> Self {
let row_number = row_number.into();
let pre = TestRowBuilder::new(row_number).with_values(pre_values).build();
let post = TestRowBuilder::new(row_number).with_values(post_values).build();
self.update(pre, post)
}
pub fn remove(mut self, row: Row) -> Self {
self.diffs.push(Diff::Remove {
pre: Columns::from_row(&row),
});
self
}
pub fn remove_row(self, row_number: impl Into<RowNumber>, values: Vec<Value>) -> Self {
let row = TestRowBuilder::new(row_number).with_values(values).build();
self.remove(row)
}
pub fn build(self) -> Change {
Change {
origin: self.origin,
diffs: self.diffs,
version: self.version,
changed_at: self.changed_at,
}
}
}
pub struct TestLayoutBuilder {
fields: Vec<RowShapeField>,
}
impl Default for TestLayoutBuilder {
fn default() -> Self {
Self::new()
}
}
impl TestLayoutBuilder {
pub fn new() -> Self {
Self {
fields: Vec::new(),
}
}
pub fn add_type(mut self, ty: Type) -> Self {
let field_name = format!("field{}", self.fields.len());
self.fields.push(RowShapeField::unconstrained(field_name, ty));
self
}
pub fn add_field(mut self, name: impl Into<String>, ty: Type) -> Self {
self.fields.push(RowShapeField::unconstrained(name, ty));
self
}
pub fn build(self) -> RowShape {
RowShape::new(self.fields)
}
pub fn build_named(self) -> RowShape {
self.build()
}
}
pub mod helpers {
use reifydb_core::{encoded::shape::RowShape, interface::change::Change, row::Row};
use reifydb_type::value::{row_number::RowNumber, r#type::Type};
use super::*;
pub fn counter_layout() -> RowShape {
TestLayoutBuilder::new().add_type(Type::Int8).build()
}
pub fn key_value_layout() -> RowShape {
TestLayoutBuilder::new().add_type(Type::Utf8).add_type(Type::Int8).build()
}
pub fn named_key_value_layout() -> RowShape {
TestLayoutBuilder::new().add_field("key", Type::Utf8).add_field("value", Type::Int8).build_named()
}
pub fn int_row(row_number: impl Into<RowNumber>, value: i8) -> Row {
TestRowBuilder::new(row_number).with_values(vec![Value::Int8(value as i64)]).build()
}
pub fn key_value_row(row_number: impl Into<RowNumber>, key: &str, value: i8) -> Row {
TestRowBuilder::new(row_number)
.with_values(vec![Value::Utf8(key.into()), Value::Int8(value as i64)])
.build()
}
pub fn empty_change() -> Change {
TestChangeBuilder::new().build()
}
pub fn insert_change(row: Row) -> Change {
TestChangeBuilder::new().insert(row).build()
}
pub fn batch_insert_change(rows: Vec<Row>) -> Change {
let mut builder = TestChangeBuilder::new();
for row in rows {
builder = builder.insert(row);
}
builder.build()
}
}
#[cfg(test)]
pub mod tests {
use reifydb_core::{
common::CommitVersion,
interface::{catalog::shape::ShapeId, change::ChangeOrigin},
};
use reifydb_type::value::{row_number::RowNumber, r#type::Type};
use super::{helpers::*, *};
#[test]
fn test_row_builder() {
let row = TestRowBuilder::new(42)
.add_value(Value::Int8(10i64))
.add_value(Value::Utf8("test".into()))
.build();
assert_eq!(row.number, RowNumber(42));
assert_eq!(row.shape.field_count(), 2);
}
#[test]
fn test_flow_change_builder() {
let change = TestChangeBuilder::new()
.changed_by_shape(ShapeId::table(100))
.with_version(CommitVersion(5))
.insert_row(1, vec![Value::Int8(42i64)])
.update_row(2, vec![Value::Int8(10i64)], vec![Value::Int8(20i64)])
.remove_row(3, vec![Value::Int8(30i64)])
.build();
assert_eq!(change.version, CommitVersion(5));
assert_eq!(change.diffs.len(), 3);
match &change.origin {
ChangeOrigin::Shape(shape) => {
assert_eq!(*shape, ShapeId::table(100));
}
_ => panic!("Expected external origin"),
}
}
#[test]
fn test_layout_builder() {
let unnamed = TestLayoutBuilder::new().add_type(Type::Int8).add_type(Type::Utf8).build();
assert_eq!(unnamed.field_count(), 2);
let named = TestLayoutBuilder::new()
.add_field("count", Type::Int8)
.add_field("name", Type::Utf8)
.build_named();
assert_eq!(named.field_count(), 2);
assert_eq!(named.get_field_name(0).unwrap(), "count");
assert_eq!(named.get_field_name(1).unwrap(), "name");
}
#[test]
fn test_helpers() {
let row = int_row(1, 42);
assert_eq!(row.number, RowNumber(1));
let kv_row = key_value_row(2, "test", 100);
assert_eq!(kv_row.number, RowNumber(2));
let change = insert_change(row.clone());
assert_eq!(change.diffs.len(), 1);
let batch = batch_insert_change(vec![row.clone(), kv_row.clone()]);
assert_eq!(batch.diffs.len(), 2);
}
}