1use 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
16pub struct TestRowBuilder {
18 row_number: RowNumber,
19 values: Vec<Value>,
20 shape: Option<RowShape>,
21}
22
23impl TestRowBuilder {
24 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 pub fn with_values(mut self, values: Vec<Value>) -> Self {
35 self.values = values;
36 self
37 }
38
39 pub fn add_value(mut self, value: Value) -> Self {
41 self.values.push(value);
42 self
43 }
44
45 pub fn with_shape(mut self, shape: RowShape) -> Self {
47 self.shape = Some(shape);
48 self
49 }
50
51 pub fn build(self) -> Row {
53 let shape = if let Some(shape) = self.shape {
55 shape
56 } else {
57 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
78pub 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 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 pub fn changed_by_shape(mut self, shape: ShapeId) -> Self {
105 self.origin = ChangeOrigin::Shape(shape);
106 self
107 }
108
109 pub fn changed_by_node(mut self, node: FlowNodeId) -> Self {
111 self.origin = ChangeOrigin::Flow(node);
112 self
113 }
114
115 pub fn with_version(mut self, version: CommitVersion) -> Self {
117 self.version = version;
118 self
119 }
120
121 pub fn with_changed_at(mut self, changed_at: DateTime) -> Self {
123 self.changed_at = changed_at;
124 self
125 }
126
127 pub fn insert(mut self, row: Row) -> Self {
129 self.diffs.push(Diff::insert(Columns::from_row(&row)));
130 self
131 }
132
133 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 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 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 pub fn remove(mut self, row: Row) -> Self {
160 self.diffs.push(Diff::remove(Columns::from_row(&row)));
161 self
162 }
163
164 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 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
181pub 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 pub fn new() -> Self {
195 Self {
196 fields: Vec::new(),
197 }
198 }
199
200 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 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 pub fn build(self) -> RowShape {
215 RowShape::new(self.fields)
216 }
217
218 pub fn build_named(self) -> RowShape {
220 self.build()
221 }
222}
223
224pub 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 pub fn counter_layout() -> RowShape {
233 TestLayoutBuilder::new().add_type(Type::Int8).build()
234 }
235
236 pub fn key_value_layout() -> RowShape {
238 TestLayoutBuilder::new().add_type(Type::Utf8).add_type(Type::Int8).build()
239 }
240
241 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 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 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 pub fn empty_change() -> Change {
260 TestChangeBuilder::new().build()
261 }
262
263 pub fn insert_change(row: Row) -> Change {
265 TestChangeBuilder::new().insert(row).build()
266 }
267
268 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}