Skip to main content

reifydb_sdk/testing/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::{
5	collections::HashMap,
6	sync::{Arc, Mutex},
7};
8
9use reifydb_core::{
10	common::CommitVersion,
11	encoded::{key::EncodedKey, row::EncodedRow},
12};
13use reifydb_type::util::cowvec::CowVec;
14
15#[derive(Clone)]
16pub struct TestContext {
17	state_store: Arc<Mutex<HashMap<EncodedKey, EncodedRow>>>,
18	version: CommitVersion,
19	logs: Arc<Mutex<Vec<String>>>,
20}
21
22impl Default for TestContext {
23	fn default() -> Self {
24		Self::new(CommitVersion(1))
25	}
26}
27
28impl TestContext {
29	pub fn new(version: CommitVersion) -> Self {
30		Self {
31			state_store: Arc::new(Mutex::new(HashMap::new())),
32			version,
33			logs: Arc::new(Mutex::new(Vec::new())),
34		}
35	}
36
37	pub fn state_store(&self) -> &Arc<Mutex<HashMap<EncodedKey, EncodedRow>>> {
38		&self.state_store
39	}
40
41	pub fn logs(&self) -> Vec<String> {
42		self.logs.lock().unwrap().clone()
43	}
44
45	pub fn clear_logs(&self) {
46		self.logs.lock().unwrap().clear();
47	}
48
49	pub fn version(&self) -> CommitVersion {
50		self.version
51	}
52
53	pub fn set_version(&mut self, version: CommitVersion) {
54		self.version = version;
55	}
56
57	pub fn get_state(&self, key: &EncodedKey) -> Option<Vec<u8>> {
58		self.state_store.lock().unwrap().get(key).map(|v| v.0.to_vec())
59	}
60
61	pub fn set_state(&self, key: EncodedKey, value: Vec<u8>) {
62		self.state_store.lock().unwrap().insert(key, EncodedRow(CowVec::new(value)));
63	}
64
65	pub fn remove_state(&self, key: &EncodedKey) -> Option<Vec<u8>> {
66		self.state_store.lock().unwrap().remove(key).map(|v| v.0.to_vec())
67	}
68
69	pub fn has_state(&self, key: &EncodedKey) -> bool {
70		self.state_store.lock().unwrap().contains_key(key)
71	}
72
73	pub fn state_count(&self) -> usize {
74		self.state_store.lock().unwrap().len()
75	}
76
77	pub fn clear_state(&self) {
78		self.state_store.lock().unwrap().clear();
79	}
80
81	pub fn state_keys(&self) -> Vec<EncodedKey> {
82		self.state_store.lock().unwrap().keys().cloned().collect()
83	}
84}
85
86#[cfg(test)]
87pub mod tests {
88	use super::*;
89	use crate::testing::helpers::encode_key;
90
91	#[test]
92	fn test_context_state_operations() {
93		let ctx = TestContext::default();
94		let key = encode_key("test_key");
95		let value = vec![1, 2, 3];
96
97		// Test set and get
98		ctx.set_state(key.clone(), value.clone());
99		assert_eq!(ctx.get_state(&key), Some(value.clone()));
100		assert!(ctx.has_state(&key));
101
102		// Test remove
103		let removed = ctx.remove_state(&key);
104		assert_eq!(removed, Some(value));
105		assert!(!ctx.has_state(&key));
106		assert_eq!(ctx.get_state(&key), None);
107	}
108
109	#[test]
110	fn test_context_logs() {
111		let ctx = TestContext::default();
112
113		// Simulate logging (would be done through callbacks in real usage)
114		ctx.logs.lock().unwrap().push("Log 1".to_string());
115		ctx.logs.lock().unwrap().push("Log 2".to_string());
116
117		let logs = ctx.logs();
118		assert_eq!(logs.len(), 2);
119		assert_eq!(logs[0], "Log 1");
120		assert_eq!(logs[1], "Log 2");
121
122		ctx.clear_logs();
123		assert_eq!(ctx.logs().len(), 0);
124	}
125
126	#[test]
127	fn test_context_state_inspection() {
128		let ctx = TestContext::default();
129
130		ctx.set_state(encode_key("key1"), vec![1]);
131		ctx.set_state(encode_key("key2"), vec![2]);
132		ctx.set_state(encode_key("key3"), vec![3]);
133
134		assert_eq!(ctx.state_count(), 3);
135
136		let keys = ctx.state_keys();
137		assert_eq!(keys.len(), 3);
138
139		ctx.clear_state();
140		assert_eq!(ctx.state_count(), 0);
141	}
142}