openws-server 0.1.5

WebSocket server with room-based messaging, private messages, HTTP broadcast, and configurable limits
use std::collections::HashMap;
use tokio::sync::mpsc;
use uuid::Uuid;

use crate::message::ServerMessage;

#[derive(Debug)]
pub struct Room {
    pub name: String,
    pub clients: HashMap<Uuid, ConnectedClient>,
}

#[derive(Debug)]
pub struct ConnectedClient {
    pub id: Uuid,
    pub username: String,
    pub sender: mpsc::Sender<String>,
}

impl Room {
    pub fn new(name: String) -> Self {
        Self {
            name,
            clients: HashMap::new(),
        }
    }

    pub fn broadcast(&mut self, msg: &ServerMessage, exclude_id: Option<Uuid>) {
        let text = serde_json::to_string(msg).expect("failed to serialize message");
        self.clients.retain(|id, client| {
            if Some(*id) == exclude_id {
                return true;
            }
            client.sender.try_send(text.clone()).is_ok()
        });
    }

    pub fn find_client_by_username(&self, username: &str) -> Option<&ConnectedClient> {
        self.clients.values().find(|c| c.username == username)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_client(id: Uuid, username: &str) -> (ConnectedClient, mpsc::Receiver<String>) {
        let (tx, rx) = mpsc::channel(10);
        let client = ConnectedClient {
            id,
            username: username.into(),
            sender: tx,
        };
        (client, rx)
    }

    #[test]
    fn test_room_new() {
        let room = Room::new("general".into());
        assert_eq!(room.name, "general");
        assert!(room.clients.is_empty());
    }

    #[test]
    fn test_room_add_client() {
        let mut room = Room::new("general".into());
        let (client, _) = make_client(Uuid::new_v4(), "alice");
        room.clients.insert(client.id, client);
        assert_eq!(room.clients.len(), 1);
    }

    #[test]
    fn test_room_broadcast() {
        let mut room = Room::new("general".into());
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let (client1, mut rx1) = make_client(id1, "alice");
        let (client2, mut rx2) = make_client(id2, "bob");
        room.clients.insert(id1, client1);
        room.clients.insert(id2, client2);

        let msg = ServerMessage::Message {
            room: "general".into(),
            from: "alice".into(),
            text: "hi".into(),
            at: 100,
        };

        room.broadcast(&msg, None);
        assert!(rx1.try_recv().is_ok());
        assert!(rx2.try_recv().is_ok());
    }

    #[test]
    fn test_room_broadcast_exclude() {
        let mut room = Room::new("general".into());
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let (client1, mut rx1) = make_client(id1, "alice");
        let (client2, mut rx2) = make_client(id2, "bob");
        room.clients.insert(id1, client1);
        room.clients.insert(id2, client2);

        let msg = ServerMessage::Message {
            room: "general".into(),
            from: "alice".into(),
            text: "hi".into(),
            at: 100,
        };

        room.broadcast(&msg, Some(id1));
        assert!(rx1.try_recv().is_err());
        assert!(rx2.try_recv().is_ok());
    }

    #[test]
    fn test_room_remove_empty() {
        let mut room = Room::new("general".into());
        let id = Uuid::new_v4();
        let (client, _) = make_client(id, "alice");
        room.clients.insert(id, client);
        room.clients.remove(&id);
        assert!(room.clients.is_empty());
    }

    #[test]
    fn test_find_client_by_username() {
        let mut room = Room::new("general".into());
        let id = Uuid::new_v4();
        let (client, _) = make_client(id, "alice");
        room.clients.insert(id, client);
        let found = room.find_client_by_username("alice");
        assert!(found.is_some());
        assert_eq!(found.unwrap().id, id);
        assert!(room.find_client_by_username("bob").is_none());
    }
}