1use std::collections::HashMap;
5
6use reifydb_core::encoded::{key::EncodedKey, row::EncodedRow, shape::RowShape};
7use reifydb_type::value::Value;
8
9use super::helpers::get_values;
10
11#[derive(Debug, Clone, Default)]
12pub struct TestStateStore {
13 data: HashMap<EncodedKey, EncodedRow>,
14}
15
16impl TestStateStore {
17 pub fn new() -> Self {
18 Self {
19 data: HashMap::new(),
20 }
21 }
22
23 pub fn get(&self, key: &EncodedKey) -> Option<&EncodedRow> {
24 self.data.get(key)
25 }
26
27 pub fn set(&mut self, key: EncodedKey, value: EncodedRow) {
28 self.data.insert(key, value);
29 }
30
31 pub fn remove(&mut self, key: &EncodedKey) -> Option<EncodedRow> {
32 self.data.remove(key)
33 }
34
35 pub fn contains(&self, key: &EncodedKey) -> bool {
36 self.data.contains_key(key)
37 }
38
39 pub fn len(&self) -> usize {
40 self.data.len()
41 }
42
43 pub fn is_empty(&self) -> bool {
44 self.data.is_empty()
45 }
46
47 pub fn clear(&mut self) {
48 self.data.clear();
49 }
50
51 pub fn keys(&self) -> Vec<&EncodedKey> {
52 self.data.keys().collect()
53 }
54
55 pub fn entries(&self) -> Vec<(&EncodedKey, &EncodedRow)> {
56 self.data.iter().collect()
57 }
58
59 pub fn decode_value(&self, key: &EncodedKey, shape: &RowShape) -> Option<Vec<Value>> {
60 self.get(key).map(|encoded| get_values(shape, encoded))
61 }
62
63 pub fn decode_named_value(&self, key: &EncodedKey, shape: &RowShape) -> Option<HashMap<String, Value>> {
64 self.get(key).map(|encoded| {
65 let values = get_values(shape, encoded);
66 shape.field_names().map(|n| n.to_string()).zip(values).collect()
67 })
68 }
69
70 pub fn set_value(&mut self, key: EncodedKey, values: &[Value], shape: &RowShape) {
71 let mut encoded = shape.allocate();
72 shape.set_values(&mut encoded, values);
73 self.set(key, encoded);
74 }
75
76 pub fn set_named_value(&mut self, key: EncodedKey, values: &HashMap<String, Value>, shape: &RowShape) {
77 let mut encoded = shape.allocate();
78
79 let ordered_values: Vec<Value> =
81 shape.field_names().map(|name| values.get(name).cloned().unwrap_or(Value::none())).collect();
82
83 shape.set_values(&mut encoded, &ordered_values);
84 self.set(key, encoded);
85 }
86
87 pub fn snapshot(&self) -> HashMap<EncodedKey, EncodedRow> {
88 self.data.clone()
89 }
90
91 pub fn restore(&mut self, snapshot: HashMap<EncodedKey, EncodedRow>) {
92 self.data = snapshot;
93 }
94
95 pub fn assert_value(&self, key: &EncodedKey, expected: &[Value], shape: &RowShape) {
96 let actual =
97 self.decode_value(key, shape).unwrap_or_else(|| panic!("Key {:?} not found in state", key));
98 assert_eq!(actual, expected, "State value mismatch for key {:?}", key);
99 }
100
101 pub fn assert_exists(&self, key: &EncodedKey) {
102 assert!(self.contains(key), "Expected key {:?} to exist in state", key);
103 }
104
105 pub fn assert_not_exists(&self, key: &EncodedKey) {
106 assert!(!self.contains(key), "Expected key {:?} to not exist in state", key);
107 }
108
109 pub fn assert_count(&self, expected: usize) {
110 assert_eq!(self.len(), expected, "Expected {} entries in state, found {}", expected, self.len());
111 }
112}
113
114#[cfg(test)]
115pub mod tests {
116 use reifydb_core::encoded::{
117 row::EncodedRow,
118 shape::{RowShape, RowShapeField},
119 };
120 use reifydb_type::{util::cowvec::CowVec, value::r#type::Type};
121
122 use super::*;
123 use crate::testing::helpers::encode_key;
124
125 #[test]
126 fn test_state_store_basic_operations() {
127 let mut store = TestStateStore::new();
128 let key = encode_key("test_key");
129 let value = EncodedRow(CowVec::new(vec![1, 2, 3, 4]));
130
131 assert!(store.is_empty());
132
133 store.set(key.clone(), value.clone());
134 assert_eq!(store.get(&key), Some(&value));
135 assert!(store.contains(&key));
136 assert_eq!(store.len(), 1);
137
138 let removed = store.remove(&key);
139 assert_eq!(removed, Some(value));
140 assert!(store.is_empty());
141 }
142
143 #[test]
144 fn test_state_store_with_shape() {
145 let mut store = TestStateStore::new();
146 let shape = RowShape::testing(&[Type::Int8, Type::Utf8]);
147 let key = encode_key("test_key");
148 let values = vec![Value::Int8(42i64), Value::Utf8("hello".into())];
149
150 store.set_value(key.clone(), &values, &shape);
151
152 let decoded = store.decode_value(&key, &shape).unwrap();
153 assert_eq!(decoded, values);
154 }
155
156 #[test]
157 fn test_state_store_with_named_shape() {
158 let mut store = TestStateStore::new();
159 let shape = RowShape::new(vec![
160 RowShapeField::unconstrained("count", Type::Int8),
161 RowShapeField::unconstrained("name", Type::Utf8),
162 ]);
163 let key = encode_key("test_key");
164
165 let mut values = HashMap::new();
166 values.insert("count".to_string(), Value::Int8(10i64));
167 values.insert("name".to_string(), Value::Utf8("test".into()));
168
169 store.set_named_value(key.clone(), &values, &shape);
170
171 let decoded = store.decode_named_value(&key, &shape).unwrap();
172 assert_eq!(decoded, values);
173 }
174
175 #[test]
176 fn test_state_store_snapshot_and_restore() {
177 let mut store = TestStateStore::new();
178 let key1 = encode_key("key1");
179 let key2 = encode_key("key2");
180
181 store.set(key1.clone(), EncodedRow(CowVec::new(vec![1])));
182 store.set(key2.clone(), EncodedRow(CowVec::new(vec![2])));
183
184 let snapshot = store.snapshot();
185 assert_eq!(snapshot.len(), 2);
186
187 store.clear();
188 assert!(store.is_empty());
189
190 store.restore(snapshot);
191 assert_eq!(store.len(), 2);
192 assert_eq!(store.get(&key1), Some(&EncodedRow(CowVec::new(vec![1]))));
193 assert_eq!(store.get(&key2), Some(&EncodedRow(CowVec::new(vec![2]))));
194 }
195
196 #[test]
197 fn test_state_store_assertions() {
198 let mut store = TestStateStore::new();
199 let shape = RowShape::testing(&[Type::Int8]);
200 let key = encode_key("test_key");
201 let values = vec![Value::Int8(100i64)];
202
203 store.set_value(key.clone(), &values, &shape);
204
205 store.assert_exists(&key);
206 store.assert_value(&key, &values, &shape);
207 store.assert_count(1);
208
209 let missing_key = encode_key("missing");
210 store.assert_not_exists(&missing_key);
211 }
212}