reifydb_sdk/testing/

builders.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2025 ReifyDB

use reifydb_core::{
	common::CommitVersion,
	encoded::shape::{RowShape, RowShapeField},
	interface::{
		catalog::{flow::FlowNodeId, id::TableId, shape::ShapeId},
		change::{Change, ChangeOrigin, Diff, Diffs},
	},
	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: Diffs,
	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: Diffs::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(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(Columns::from_row(&pre), 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(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);
	}
}