reifydb_sdk/testing/
context.rs1use 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 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 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 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}