reifydb_sub_flow/testing/
harness.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{
5	collections::HashMap,
6	marker::PhantomData,
7	mem,
8	ops::{Bound, Index},
9};
10
11use reifydb_catalog::catalog::Catalog;
12use reifydb_codec::{encoded::row::EncodedRow, key::encoded::EncodedKey};
13use reifydb_core::{
14	actors::pending::Pending,
15	common::CommitVersion,
16	interface::{catalog::flow::FlowNodeId, change::Change},
17	row::Row,
18};
19use reifydb_engine::test_harness::TestEngine;
20use reifydb_runtime::context::clock::{Clock, MockClock};
21use reifydb_sdk::{
22	config::Config,
23	operator::{
24		OperatorLogic, OperatorMetadata,
25		context::{OperatorContext, StateApi, StoreApi},
26	},
27	testing::builders::TestChangeBuilder,
28};
29use reifydb_transaction::interceptor::interceptors::Interceptors;
30use reifydb_value::{Result, value::Value};
31use serde::de::DeserializeOwned;
32
33use crate::{
34	operator::{
35		Operator,
36		context::native::NativeOperatorContext,
37		native::{FlowNativeBridge, NativeBridgedOperator, NativeOperatorAdapter},
38	},
39	transaction::{DeferredParams, FlowTransaction, allocators::FlowAllocators},
40};
41
42pub struct NativeOperatorHarness<C: OperatorLogic + OperatorMetadata + 'static> {
43	engine: TestEngine,
44	operator: NativeBridgedOperator,
45	node_id: FlowNodeId,
46	version: u64,
47	pending: Pending,
48	allocators: FlowAllocators,
49	current: Option<FlowTransaction>,
50	history: Vec<Change>,
51	_phantom: PhantomData<C>,
52}
53
54impl<C: OperatorLogic + OperatorMetadata + 'static> NativeOperatorHarness<C> {
55	pub fn builder() -> NativeOperatorHarnessBuilder<C> {
56		NativeOperatorHarnessBuilder::new()
57	}
58
59	fn begin_txn(&mut self) -> FlowTransaction {
60		let query = self.engine.multi().begin_query().expect("begin_query");
61		let state_query = self.engine.multi().begin_query().expect("begin_query");
62		FlowTransaction::deferred_from_parts(DeferredParams {
63			version: CommitVersion(self.version),
64			pending: mem::take(&mut self.pending),
65			query,
66			state_query,
67			dictionary_query: None,
68			single: self.engine.inner().single().clone(),
69			catalog: Catalog::testing(),
70			interceptors: Interceptors::new(),
71			clock: Clock::Mock(MockClock::from_millis(1000)),
72			allocators: self.allocators.clone(),
73		})
74	}
75
76	fn end_txn(&mut self, mut txn: FlowTransaction) {
77		self.pending = txn.take_pending();
78		self.version += 1;
79	}
80
81	pub fn apply(&mut self, input: Change) -> Result<Change> {
82		let mut txn = self.begin_txn();
83		let output = self.operator.apply(&mut txn, input)?;
84		txn.flush_operator_states()?;
85		self.end_txn(txn);
86		self.history.push(output.clone());
87		Ok(output)
88	}
89
90	pub fn apply_without_flush(&mut self, input: Change) -> Result<Change> {
91		let mut txn = self.begin_txn();
92		let output = self.operator.apply(&mut txn, input)?;
93		self.current = Some(txn);
94		self.history.push(output.clone());
95		Ok(output)
96	}
97
98	pub fn flush(&mut self) -> Result<()> {
99		let mut txn = match self.current.take() {
100			Some(txn) => txn,
101			None => self.begin_txn(),
102		};
103		txn.flush_operator_states()?;
104		self.end_txn(txn);
105		Ok(())
106	}
107
108	pub fn state_value<V: DeserializeOwned>(&mut self, key: &EncodedKey) -> Option<V> {
109		let node = self.node_id;
110		if let Some(txn) = self.current.as_mut() {
111			let mut bridge = FlowNativeBridge::new(txn, node);
112			let mut ctx = NativeOperatorContext::new(&mut bridge, node);
113			return ctx.state().get::<V>(key).expect("state get");
114		}
115		let mut txn = self.begin_txn();
116		let value = {
117			let mut bridge = FlowNativeBridge::new(&mut txn, node);
118			let mut ctx = NativeOperatorContext::new(&mut bridge, node);
119			ctx.state().get::<V>(key).expect("state get")
120		};
121		self.end_txn(txn);
122		value
123	}
124
125	pub fn seed_store(&mut self, rows: &[(EncodedKey, EncodedRow)]) {
126		let keys: Vec<EncodedKey> = rows.iter().map(|(k, _)| k.clone()).collect();
127		let values: Vec<EncodedRow> = rows.iter().map(|(_, v)| v.clone()).collect();
128		let mut txn = self.begin_txn();
129		txn.set_batch(&keys, &values).expect("seed_store set_batch");
130		self.end_txn(txn);
131	}
132
133	pub fn store_range(
134		&mut self,
135		start: Bound<&EncodedKey>,
136		end: Bound<&EncodedKey>,
137	) -> Vec<(EncodedKey, EncodedRow)> {
138		let node = self.node_id;
139		let mut txn = self.begin_txn();
140		let rows = {
141			let mut bridge = FlowNativeBridge::new(&mut txn, node);
142			let mut ctx = NativeOperatorContext::new(&mut bridge, node);
143			ctx.store().range(start, end).expect("store range")
144		};
145		self.end_txn(txn);
146		rows
147	}
148
149	pub fn insert(&mut self, row: Row) -> &mut Self {
150		let change = TestChangeBuilder::new().insert(row).build();
151		self.apply(change).expect("insert failed");
152		self
153	}
154
155	pub fn update(&mut self, pre: Row, post: Row) -> &mut Self {
156		let change = TestChangeBuilder::new().update(pre, post).build();
157		self.apply(change).expect("update failed");
158		self
159	}
160
161	pub fn remove(&mut self, row: Row) -> &mut Self {
162		let change = TestChangeBuilder::new().remove(row).build();
163		self.apply(change).expect("remove failed");
164		self
165	}
166
167	pub fn history_len(&self) -> usize {
168		self.history.len()
169	}
170
171	pub fn last_change(&self) -> Option<&Change> {
172		self.history.last()
173	}
174
175	pub fn clear_history(&mut self) {
176		self.history.clear();
177	}
178
179	pub fn node_id(&self) -> FlowNodeId {
180		self.node_id
181	}
182}
183
184impl<C: OperatorLogic + OperatorMetadata + 'static> Index<usize> for NativeOperatorHarness<C> {
185	type Output = Change;
186
187	fn index(&self, index: usize) -> &Self::Output {
188		&self.history[index]
189	}
190}
191
192pub struct NativeOperatorHarnessBuilder<C> {
193	config: HashMap<String, Value>,
194	node_id: FlowNodeId,
195	version: CommitVersion,
196	_phantom: PhantomData<C>,
197}
198
199impl<C: OperatorLogic + OperatorMetadata + 'static> Default for NativeOperatorHarnessBuilder<C> {
200	fn default() -> Self {
201		Self::new()
202	}
203}
204
205impl<C: OperatorLogic + OperatorMetadata + 'static> NativeOperatorHarnessBuilder<C> {
206	pub fn new() -> Self {
207		Self {
208			config: HashMap::new(),
209			node_id: FlowNodeId(1),
210			version: CommitVersion(1),
211			_phantom: PhantomData,
212		}
213	}
214
215	pub fn with_config<I, K>(mut self, config: I) -> Self
216	where
217		I: IntoIterator<Item = (K, Value)>,
218		K: Into<String>,
219	{
220		self.config = config.into_iter().map(|(k, v)| (k.into(), v)).collect();
221		self
222	}
223
224	pub fn add_config(mut self, key: impl Into<String>, value: Value) -> Self {
225		self.config.insert(key.into(), value);
226		self
227	}
228
229	pub fn with_node_id(mut self, node_id: FlowNodeId) -> Self {
230		self.node_id = node_id;
231		self
232	}
233
234	pub fn with_version(mut self, version: CommitVersion) -> Self {
235		self.version = version;
236		self
237	}
238
239	pub fn build(self) -> Result<NativeOperatorHarness<C>> {
240		let engine = TestEngine::new();
241		let core = C::create(
242			self.node_id,
243			&Config::new(<C as OperatorMetadata>::NAME, self.config.clone().into_iter().collect()),
244		)?;
245		let capabilities = <C as OperatorMetadata>::CAPABILITIES;
246		let adapter = NativeOperatorAdapter::new(core, self.node_id, capabilities);
247		let operator = NativeBridgedOperator::new(Box::new(adapter), self.node_id, capabilities);
248
249		Ok(NativeOperatorHarness {
250			engine,
251			operator,
252			node_id: self.node_id,
253			version: self.version.0,
254			pending: Pending::new(),
255			allocators: FlowAllocators::new(),
256			current: None,
257			history: Vec::new(),
258			_phantom: PhantomData,
259		})
260	}
261}
reifydb_sub_flow/testing/harness.rs

reifydb_sub_flow/testing/
harness.rs
