1use serde_json::{json, Value};
2use std::collections::HashMap;
3use std::sync::{Arc, Mutex};
4use tokio::sync::mpsc;
5use tokio_tungstenite::tungstenite::protocol::Message;
6
7use crate::error::{ObnizError, ObnizResult};
8use crate::obniz::{CallbackType, ObnizCommand};
9
10#[derive(Debug, Clone)]
12pub struct MockMessage {
13 pub request: Value,
14 pub response: Value,
15 pub delay_ms: Option<u64>,
16}
17
18#[derive(Debug, Clone)]
20pub struct MockConfig {
21 pub device_id: String,
22 pub should_fail_connection: bool,
23 pub should_timeout: bool,
24 pub default_delay_ms: u64,
25}
26
27impl Default for MockConfig {
28 fn default() -> Self {
29 Self {
30 device_id: "mock-device".to_string(),
31 should_fail_connection: false,
32 should_timeout: false,
33 default_delay_ms: 10,
34 }
35 }
36}
37
38#[derive(Debug)]
40pub struct MockWebSocketServer {
41 config: MockConfig,
42 message_handlers: Arc<Mutex<HashMap<String, MockMessage>>>,
43 sent_messages: Arc<Mutex<Vec<Value>>>,
44 callbacks: Arc<Mutex<HashMap<String, CallbackType>>>,
45}
46
47impl MockWebSocketServer {
48 pub fn new(config: MockConfig) -> Self {
49 Self {
50 config,
51 message_handlers: Arc::new(Mutex::new(HashMap::new())),
52 sent_messages: Arc::new(Mutex::new(Vec::new())),
53 callbacks: Arc::new(Mutex::new(HashMap::new())),
54 }
55 }
56
57 pub fn add_response(&self, request_key: &str, response: Value) {
59 let mock_msg = MockMessage {
60 request: json!({}),
61 response,
62 delay_ms: Some(self.config.default_delay_ms),
63 };
64 self.message_handlers
65 .lock()
66 .unwrap()
67 .insert(request_key.to_string(), mock_msg);
68 }
69
70 pub fn add_response_with_delay(&self, request_key: &str, response: Value, delay_ms: u64) {
72 let mock_msg = MockMessage {
73 request: json!({}),
74 response,
75 delay_ms: Some(delay_ms),
76 };
77 self.message_handlers
78 .lock()
79 .unwrap()
80 .insert(request_key.to_string(), mock_msg);
81 }
82
83 pub fn get_sent_messages(&self) -> Vec<Value> {
85 self.sent_messages.lock().unwrap().clone()
86 }
87
88 pub fn clear_sent_messages(&self) {
90 self.sent_messages.lock().unwrap().clear();
91 }
92
93 pub async fn process_message(&self, message: Message) -> ObnizResult<Option<Value>> {
95 if self.config.should_timeout {
96 tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
97 }
98
99 let text = message
100 .to_text()
101 .map_err(|_| ObnizError::Generic("Invalid message".to_string()))?;
102 let request: Value =
103 serde_json::from_str(text).map_err(|e| ObnizError::JsonParse(e.to_string()))?;
104
105 self.sent_messages.lock().unwrap().push(request.clone());
107
108 let response = self.find_mock_response(&request);
110
111 if let Some(mock_msg) = response {
112 if let Some(delay) = mock_msg.delay_ms {
113 tokio::time::sleep(tokio::time::Duration::from_millis(delay)).await;
114 }
115 Ok(Some(mock_msg.response))
116 } else {
117 Ok(Some(self.generate_default_response(&request)))
119 }
120 }
121
122 fn find_mock_response(&self, request: &Value) -> Option<MockMessage> {
123 let handlers = self.message_handlers.lock().unwrap();
124
125 if let Some(array) = request.as_array() {
127 if let Some(first_item) = array.first() {
128 if let Some(obj) = first_item.as_object() {
129 for (key, _) in obj {
130 if let Some(mock_msg) = handlers.get(key) {
131 return Some(mock_msg.clone());
132 }
133
134 if let Some(nested) = first_item.get(key) {
136 if let Some(nested_obj) = nested.as_object() {
137 for (nested_key, _) in nested_obj {
138 let compound_key = format!("{key}.{nested_key}");
139 if let Some(mock_msg) = handlers.get(&compound_key) {
140 return Some(mock_msg.clone());
141 }
142 }
143 }
144 }
145 }
146 }
147 }
148 }
149 None
150 }
151
152 fn generate_default_response(&self, request: &Value) -> Value {
153 if let Some(array) = request.as_array() {
154 if let Some(first_item) = array.first() {
155 if let Some(obj) = first_item.as_object() {
156 for (key, value) in obj {
157 match key.as_str() {
158 k if k.starts_with("io") => {
160 if value == "get" {
161 return json!([{k: false}]); } else {
163 return json!([{k: {"state": "ok"}}]);
164 }
165 }
166 k if k.starts_with("ad") => {
168 if value == "get" {
169 return json!([{k: 3.3}]); } else {
171 return json!([{k: {"state": "ok"}}]);
172 }
173 }
174 k if k.starts_with("pwm") => {
176 return json!([{k: {"state": "ok"}}]);
177 }
178 k if k.starts_with("uart") => {
180 return json!([{k: {"state": "ok"}}]);
181 }
182 "display" => {
184 return json!([{"display": {"state": "ok"}}]);
185 }
186 "system" => {
188 if let Some(system_obj) = value.as_object() {
189 if system_obj.contains_key("info") {
190 return json!([{
191 "system": {
192 "version": "3.5.0",
193 "hardware": "mock",
194 "device_type": "mock",
195 "region": "test"
196 }
197 }]);
198 }
199 }
200 return json!([{"system": {"state": "ok"}}]);
201 }
202 "switch" => {
204 if value == "get" {
205 return json!([{
206 "switch": {
207 "state": "none",
208 "action": "get"
209 }
210 }]);
211 }
212 return json!([{"switch": {"state": "ok"}}]);
213 }
214 "ws" => {
216 return json!([{
217 "ws": {
218 "ready": true,
219 "obniz": {
220 "hw": "mock",
221 "firmware": "test",
222 "connected_network": {
223 "online_at": 1640995200,
224 "wifi": {
225 "ssid": "test-wifi"
226 }
227 }
228 }
229 }
230 }]);
231 }
232 _ => {}
233 }
234 }
235 }
236 }
237 }
238
239 json!([{"status": "ok"}])
241 }
242
243 pub async fn trigger_callback(&self, key: &str, data: Value) {
245 if let Some(callback) = self.callbacks.lock().unwrap().get(key) {
246 match callback {
247 CallbackType::Persistent(callback_fn) => {
248 callback_fn(data);
249 }
250 CallbackType::OneShot(_) => {
251 }
254 }
255 }
256 }
257}
258
259#[derive(Debug)]
261pub struct MockObniz {
262 device_id: String,
263 server: Arc<MockWebSocketServer>,
264 command_sender: mpsc::UnboundedSender<ObnizCommand>,
265 #[allow(dead_code)] command_receiver: Arc<Mutex<Option<mpsc::UnboundedReceiver<ObnizCommand>>>>,
267}
268
269impl MockObniz {
270 pub fn new(config: MockConfig) -> Self {
271 let (tx, rx) = mpsc::unbounded_channel();
272 let server = Arc::new(MockWebSocketServer::new(config.clone()));
273
274 Self {
275 device_id: config.device_id,
276 server,
277 command_sender: tx,
278 command_receiver: Arc::new(Mutex::new(Some(rx))),
279 }
280 }
281
282 pub fn server(&self) -> Arc<MockWebSocketServer> {
283 self.server.clone()
284 }
285
286 pub async fn send_message(&self, message: Message) -> ObnizResult<()> {
287 self.command_sender
288 .send(ObnizCommand::Send {
289 message,
290 response_key: None,
291 })
292 .map_err(|e| ObnizError::Connection(e.to_string()))?;
293 Ok(())
294 }
295
296 pub async fn send_await_response(
297 &self,
298 message: Message,
299 _response_key: String,
300 ) -> ObnizResult<Value> {
301 if let Some(response) = self.server.process_message(message).await? {
303 Ok(response)
304 } else {
305 Err(ObnizError::Timeout)
306 }
307 }
308
309 pub fn register_callback<F>(&self, key: String, callback: F) -> ObnizResult<()>
310 where
311 F: Fn(Value) + Send + Sync + 'static,
312 {
313 self.server
314 .callbacks
315 .lock()
316 .unwrap()
317 .insert(key, CallbackType::Persistent(Box::new(callback)));
318 Ok(())
319 }
320
321 pub fn unregister_callback(&self, key: String) -> ObnizResult<()> {
322 self.server.callbacks.lock().unwrap().remove(&key);
323 Ok(())
324 }
325
326 pub fn id(&self) -> &str {
327 &self.device_id
328 }
329}
330
331pub mod responses {
333 use super::*;
334
335 pub fn io_pin_state(pin: u8, state: bool) -> Value {
336 json!([{format!("io{}", pin): state}])
337 }
338
339 pub fn ad_voltage(channel: u8, voltage: f64) -> Value {
340 json!([{format!("ad{}", channel): voltage}])
341 }
342
343 pub fn pwm_ok(channel: u8) -> Value {
344 json!([{format!("pwm{}", channel): {"state": "ok"}}])
345 }
346
347 pub fn uart_data(channel: u8, data: Vec<u8>) -> Value {
348 json!([{format!("uart{}", channel): {"data": data}}])
349 }
350
351 pub fn switch_state(state: &str, action: &str) -> Value {
352 json!([{
353 "switch": {
354 "state": state,
355 "action": action
356 }
357 }])
358 }
359
360 pub fn display_ok() -> Value {
361 json!([{"display": {"state": "ok"}}])
362 }
363
364 pub fn system_info() -> Value {
365 json!([{
366 "system": {
367 "version": "3.5.0",
368 "hardware": "mock",
369 "device_type": "mock",
370 "region": "test"
371 }
372 }])
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379
380 #[tokio::test]
381 async fn test_mock_server_creation() {
382 let config = MockConfig::default();
383 let server = MockWebSocketServer::new(config);
384
385 assert_eq!(server.config.device_id, "mock-device");
386 assert!(!server.config.should_fail_connection);
387 }
388
389 #[tokio::test]
390 async fn test_mock_responses() {
391 let config = MockConfig::default();
392 let server = MockWebSocketServer::new(config);
393
394 server.add_response("io0", responses::io_pin_state(0, true));
396
397 let request = json!([{"io0": "get"}]);
399 let message = Message::from(request.to_string());
400
401 let response = server.process_message(message).await.unwrap();
402 assert!(response.is_some());
403
404 let response_value = response.unwrap();
405 assert_eq!(response_value, responses::io_pin_state(0, true));
406 }
407
408 #[tokio::test]
409 async fn test_default_responses() {
410 let config = MockConfig::default();
411 let server = MockWebSocketServer::new(config);
412
413 let request = json!([{"io0": "get"}]);
415 let message = Message::from(request.to_string());
416
417 let response = server.process_message(message).await.unwrap();
418 assert!(response.is_some());
419
420 let response_value = response.unwrap();
422 assert_eq!(response_value, json!([{"io0": false}]));
423 }
424
425 #[tokio::test]
426 async fn test_sent_message_tracking() {
427 let config = MockConfig::default();
428 let server = MockWebSocketServer::new(config);
429
430 let request1 = json!([{"io0": true}]);
431 let request2 = json!([{"io1": false}]);
432
433 let message1 = Message::from(request1.to_string());
434 let message2 = Message::from(request2.to_string());
435
436 server.process_message(message1).await.unwrap();
437 server.process_message(message2).await.unwrap();
438
439 let sent_messages = server.get_sent_messages();
440 assert_eq!(sent_messages.len(), 2);
441 assert_eq!(sent_messages[0], request1);
442 assert_eq!(sent_messages[1], request2);
443 }
444}