Skip to main content

reifydb_sdk/testing/
builders.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::{
5	common::CommitVersion,
6	encoded::shape::{RowShape, RowShapeField},
7	interface::{
8		catalog::{flow::FlowNodeId, id::TableId, shape::ShapeId},
9		change::{Change, ChangeOrigin, Diff, Diffs},
10	},
11	row::Row,
12	value::column::columns::Columns,
13};
14use reifydb_type::value::{Value, datetime::DateTime, row_number::RowNumber, r#type::Type};
15
16/// Builder for creating test rows
17pub struct TestRowBuilder {
18	row_number: RowNumber,
19	values: Vec<Value>,
20	shape: Option<RowShape>,
21}
22
23impl TestRowBuilder {
24	/// Create a new row builder with the given row number
25	pub fn new(row_number: impl Into<RowNumber>) -> Self {
26		Self {
27			row_number: row_number.into(),
28			values: Vec::new(),
29			shape: None,
30		}
31	}
32
33	/// Set the values for the row
34	pub fn with_values(mut self, values: Vec<Value>) -> Self {
35		self.values = values;
36		self
37	}
38
39	/// Add a single value to the row
40	pub fn add_value(mut self, value: Value) -> Self {
41		self.values.push(value);
42		self
43	}
44
45	/// Set the shape for the row (inferred from values if not set)
46	pub fn with_shape(mut self, shape: RowShape) -> Self {
47		self.shape = Some(shape);
48		self
49	}
50
51	/// Build the row
52	pub fn build(self) -> Row {
53		// Use provided shape or infer from values
54		let shape = if let Some(shape) = self.shape {
55			shape
56		} else {
57			// Infer types from values and create shape
58			let fields: Vec<RowShapeField> = self
59				.values
60				.iter()
61				.enumerate()
62				.map(|(i, v)| RowShapeField::unconstrained(format!("field{}", i), v.get_type()))
63				.collect();
64			RowShape::new(fields)
65		};
66
67		let mut encoded = shape.allocate();
68		shape.set_values(&mut encoded, &self.values);
69
70		Row {
71			number: self.row_number,
72			encoded,
73			shape,
74		}
75	}
76}
77
78/// Builder for creating test flow changes
79pub struct TestChangeBuilder {
80	origin: ChangeOrigin,
81	diffs: Diffs,
82	version: CommitVersion,
83	changed_at: DateTime,
84}
85
86impl Default for TestChangeBuilder {
87	fn default() -> Self {
88		Self::new()
89	}
90}
91
92impl TestChangeBuilder {
93	/// Create a new flow change builder with default origin and version
94	pub fn new() -> Self {
95		Self {
96			origin: ChangeOrigin::Shape(ShapeId::Table(TableId(1))),
97			diffs: Diffs::new(),
98			version: CommitVersion(1),
99			changed_at: DateTime::default(),
100		}
101	}
102
103	/// Set the origin as an external source
104	pub fn changed_by_shape(mut self, shape: ShapeId) -> Self {
105		self.origin = ChangeOrigin::Shape(shape);
106		self
107	}
108
109	/// Set the origin as an internal node
110	pub fn changed_by_node(mut self, node: FlowNodeId) -> Self {
111		self.origin = ChangeOrigin::Flow(node);
112		self
113	}
114
115	/// Set the version
116	pub fn with_version(mut self, version: CommitVersion) -> Self {
117		self.version = version;
118		self
119	}
120
121	/// Set the changed_at timestamp
122	pub fn with_changed_at(mut self, changed_at: DateTime) -> Self {
123		self.changed_at = changed_at;
124		self
125	}
126
127	/// Add an insert diff
128	pub fn insert(mut self, row: Row) -> Self {
129		self.diffs.push(Diff::insert(Columns::from_row(&row)));
130		self
131	}
132
133	/// Add an insert diff with values (convenience method)
134	pub fn insert_row(self, row_number: impl Into<RowNumber>, values: Vec<Value>) -> Self {
135		let row = TestRowBuilder::new(row_number).with_values(values).build();
136		self.insert(row)
137	}
138
139	/// Add an update diff
140	pub fn update(mut self, pre: Row, post: Row) -> Self {
141		self.diffs.push(Diff::update(Columns::from_row(&pre), Columns::from_row(&post)));
142		self
143	}
144
145	/// Add an update diff with values (convenience method)
146	pub fn update_row(
147		self,
148		row_number: impl Into<RowNumber>,
149		pre_values: Vec<Value>,
150		post_values: Vec<Value>,
151	) -> Self {
152		let row_number = row_number.into();
153		let pre = TestRowBuilder::new(row_number).with_values(pre_values).build();
154		let post = TestRowBuilder::new(row_number).with_values(post_values).build();
155		self.update(pre, post)
156	}
157
158	/// Add a remove diff
159	pub fn remove(mut self, row: Row) -> Self {
160		self.diffs.push(Diff::remove(Columns::from_row(&row)));
161		self
162	}
163
164	/// Add a remove diff with values (convenience method)
165	pub fn remove_row(self, row_number: impl Into<RowNumber>, values: Vec<Value>) -> Self {
166		let row = TestRowBuilder::new(row_number).with_values(values).build();
167		self.remove(row)
168	}
169
170	/// Build the flow change
171	pub fn build(self) -> Change {
172		Change {
173			origin: self.origin,
174			diffs: self.diffs,
175			version: self.version,
176			changed_at: self.changed_at,
177		}
178	}
179}
180
181/// Builder for creating test shapes
182pub struct TestLayoutBuilder {
183	fields: Vec<RowShapeField>,
184}
185
186impl Default for TestLayoutBuilder {
187	fn default() -> Self {
188		Self::new()
189	}
190}
191
192impl TestLayoutBuilder {
193	/// Create a new shape builder
194	pub fn new() -> Self {
195		Self {
196			fields: Vec::new(),
197		}
198	}
199
200	/// Add a type to the shape with auto-generated name
201	pub fn add_type(mut self, ty: Type) -> Self {
202		let field_name = format!("field{}", self.fields.len());
203		self.fields.push(RowShapeField::unconstrained(field_name, ty));
204		self
205	}
206
207	/// Add a named field to the shape
208	pub fn add_field(mut self, name: impl Into<String>, ty: Type) -> Self {
209		self.fields.push(RowShapeField::unconstrained(name, ty));
210		self
211	}
212
213	/// Build the shape
214	pub fn build(self) -> RowShape {
215		RowShape::new(self.fields)
216	}
217
218	/// Build the shape (alias for backwards compatibility)
219	pub fn build_named(self) -> RowShape {
220		self.build()
221	}
222}
223
224/// Helper functions for common test data patterns
225pub mod helpers {
226	use reifydb_core::{encoded::shape::RowShape, interface::change::Change, row::Row};
227	use reifydb_type::value::{row_number::RowNumber, r#type::Type};
228
229	use super::*;
230
231	/// Create a simple counter shape (single int8 field)
232	pub fn counter_layout() -> RowShape {
233		TestLayoutBuilder::new().add_type(Type::Int8).build()
234	}
235
236	/// Create a key-value shape (utf8 key, int8 value)
237	pub fn key_value_layout() -> RowShape {
238		TestLayoutBuilder::new().add_type(Type::Utf8).add_type(Type::Int8).build()
239	}
240
241	/// Create a named key-value shape
242	pub fn named_key_value_layout() -> RowShape {
243		TestLayoutBuilder::new().add_field("key", Type::Utf8).add_field("value", Type::Int8).build_named()
244	}
245
246	/// Create a test row with a single int8 value
247	pub fn int_row(row_number: impl Into<RowNumber>, value: i8) -> Row {
248		TestRowBuilder::new(row_number).with_values(vec![Value::Int8(value as i64)]).build()
249	}
250
251	/// Create a test row with a UTF8 key and int8 value
252	pub fn key_value_row(row_number: impl Into<RowNumber>, key: &str, value: i8) -> Row {
253		TestRowBuilder::new(row_number)
254			.with_values(vec![Value::Utf8(key.into()), Value::Int8(value as i64)])
255			.build()
256	}
257
258	/// Create an empty flow change from a table source
259	pub fn empty_change() -> Change {
260		TestChangeBuilder::new().build()
261	}
262
263	/// Create a flow change with a single insert
264	pub fn insert_change(row: Row) -> Change {
265		TestChangeBuilder::new().insert(row).build()
266	}
267
268	/// Create a flow change with multiple inserts
269	pub fn batch_insert_change(rows: Vec<Row>) -> Change {
270		let mut builder = TestChangeBuilder::new();
271		for row in rows {
272			builder = builder.insert(row);
273		}
274		builder.build()
275	}
276}
277
278#[cfg(test)]
279pub mod tests {
280	use reifydb_core::{
281		common::CommitVersion,
282		interface::{catalog::shape::ShapeId, change::ChangeOrigin},
283	};
284	use reifydb_type::value::{row_number::RowNumber, r#type::Type};
285
286	use super::{helpers::*, *};
287
288	#[test]
289	fn test_row_builder() {
290		let row = TestRowBuilder::new(42)
291			.add_value(Value::Int8(10i64))
292			.add_value(Value::Utf8("test".into()))
293			.build();
294
295		assert_eq!(row.number, RowNumber(42));
296		assert_eq!(row.shape.field_count(), 2);
297	}
298
299	#[test]
300	fn test_flow_change_builder() {
301		let change = TestChangeBuilder::new()
302			.changed_by_shape(ShapeId::table(100))
303			.with_version(CommitVersion(5))
304			.insert_row(1, vec![Value::Int8(42i64)])
305			.update_row(2, vec![Value::Int8(10i64)], vec![Value::Int8(20i64)])
306			.remove_row(3, vec![Value::Int8(30i64)])
307			.build();
308
309		assert_eq!(change.version, CommitVersion(5));
310		assert_eq!(change.diffs.len(), 3);
311
312		match &change.origin {
313			ChangeOrigin::Shape(shape) => {
314				assert_eq!(*shape, ShapeId::table(100));
315			}
316			_ => panic!("Expected external origin"),
317		}
318	}
319
320	#[test]
321	fn test_layout_builder() {
322		let unnamed = TestLayoutBuilder::new().add_type(Type::Int8).add_type(Type::Utf8).build();
323
324		assert_eq!(unnamed.field_count(), 2);
325
326		let named = TestLayoutBuilder::new()
327			.add_field("count", Type::Int8)
328			.add_field("name", Type::Utf8)
329			.build_named();
330
331		assert_eq!(named.field_count(), 2);
332		assert_eq!(named.get_field_name(0).unwrap(), "count");
333		assert_eq!(named.get_field_name(1).unwrap(), "name");
334	}
335
336	#[test]
337	fn test_helpers() {
338		let row = int_row(1, 42);
339		assert_eq!(row.number, RowNumber(1));
340
341		let kv_row = key_value_row(2, "test", 100);
342		assert_eq!(kv_row.number, RowNumber(2));
343
344		let change = insert_change(row.clone());
345		assert_eq!(change.diffs.len(), 1);
346
347		let batch = batch_insert_change(vec![row.clone(), kv_row.clone()]);
348		assert_eq!(batch.diffs.len(), 2);
349	}
350}