use crate::data_model::{DataModel, WebSocketMessage}; use std::sync::Arc; use tokio::sync::{broadcast, RwLock}; use tracing::debug; /// Manages the application state and broadcasts updates to WebSocket clients pub struct StateManager { data: Arc>, revision: Arc>, broadcast_tx: broadcast::Sender, } impl StateManager { pub fn new() -> Self { let (broadcast_tx, _) = broadcast::channel(100); Self { data: Arc::new(RwLock::new(DataModel::new())), revision: Arc::new(RwLock::new(0)), broadcast_tx, } } /// Get the current data model and revision pub async fn get_snapshot(&self) -> (DataModel, u32) { let data = self.data.read().await.clone(); let rev = *self.revision.read().await; (data, rev) } /// Subscribe to state updates pub fn subscribe(&self) -> broadcast::Receiver { self.broadcast_tx.subscribe() } /// Update the data model and broadcast to all connected clients pub async fn update_data(&self, new_data: DataModel) { let mut data = self.data.write().await; let mut rev = self.revision.write().await; *data = new_data.clone(); *rev += 1; debug!("Data model updated to revision {}", *rev); // Broadcast full data dump to all connected clients // In the future, we can optimize this by computing and sending JSON patches let message = WebSocketMessage { rev: *rev, data: Some(new_data), patch: None, }; // Ignore errors if no receivers are connected let _ = self.broadcast_tx.send(message); } /// Get a WebSocket message with the current state pub async fn get_initial_message(&self) -> WebSocketMessage { let (data, rev) = self.get_snapshot().await; WebSocketMessage { rev, data: Some(data), patch: None, } } } impl Default for StateManager { fn default() -> Self { Self::new() } } #[cfg(test)] mod tests { use super::*; use crate::data_model::DataModel; #[test] fn test_state_manager_new() { let sm = StateManager::new(); // Should be constructible without panic let _ = sm.subscribe(); } #[test] fn test_state_manager_default() { let sm = StateManager::default(); let _ = sm.subscribe(); } #[tokio::test] async fn test_get_snapshot_initial_revision_zero() { let sm = StateManager::new(); let (data, rev) = sm.get_snapshot().await; assert_eq!(rev, 0); // DataModel::new() sets version from CARGO_PKG_VERSION assert_eq!(data.server_info.version, env!("CARGO_PKG_VERSION")); assert!(data.package_data.is_empty()); assert!(data.notifications.is_empty()); } #[tokio::test] async fn test_update_data_increments_revision() { let sm = StateManager::new(); let (_, rev0) = sm.get_snapshot().await; assert_eq!(rev0, 0); sm.update_data(DataModel::new()).await; let (_, rev1) = sm.get_snapshot().await; assert_eq!(rev1, 1); sm.update_data(DataModel::new()).await; let (_, rev2) = sm.get_snapshot().await; assert_eq!(rev2, 2); } #[tokio::test] async fn test_update_data_stores_new_data() { let sm = StateManager::new(); let mut new_data = DataModel::new(); new_data.server_info.name = Some("TestNode".to_string()); new_data.ui.theme = "light".to_string(); sm.update_data(new_data).await; let (snapshot, _) = sm.get_snapshot().await; assert_eq!(snapshot.server_info.name, Some("TestNode".to_string())); assert_eq!(snapshot.ui.theme, "light"); } #[tokio::test] async fn test_subscribe_receives_updates() { let sm = StateManager::new(); let mut rx = sm.subscribe(); let mut data = DataModel::new(); data.server_info.name = Some("BroadcastTest".to_string()); sm.update_data(data).await; let msg = rx.recv().await.unwrap(); assert_eq!(msg.rev, 1); assert!(msg.data.is_some()); let received_data = msg.data.unwrap(); assert_eq!( received_data.server_info.name, Some("BroadcastTest".to_string()) ); } #[tokio::test] async fn test_multiple_subscribers_receive_same_update() { let sm = StateManager::new(); let mut rx1 = sm.subscribe(); let mut rx2 = sm.subscribe(); sm.update_data(DataModel::new()).await; let msg1 = rx1.recv().await.unwrap(); let msg2 = rx2.recv().await.unwrap(); assert_eq!(msg1.rev, msg2.rev); assert_eq!(msg1.rev, 1); } #[tokio::test] async fn test_get_initial_message() { let sm = StateManager::new(); let mut data = DataModel::new(); data.server_info.name = Some("InitMsg".to_string()); sm.update_data(data).await; let msg = sm.get_initial_message().await; assert_eq!(msg.rev, 1); assert!(msg.data.is_some()); assert!(msg.patch.is_none()); assert_eq!( msg.data.unwrap().server_info.name, Some("InitMsg".to_string()) ); } #[tokio::test] async fn test_update_without_subscribers_does_not_panic() { let sm = StateManager::new(); // No subscribers — send should silently succeed sm.update_data(DataModel::new()).await; let (_, rev) = sm.get_snapshot().await; assert_eq!(rev, 1); } }