mini rust message broker

broker.rs

use std::collections::{HashMap, VecDeque};

#[derive(Debug, Clone, PartialEq)]
pub enum Message {
    Text(String),
    Number(i64),
    Binary(Vec<u8>),
    Batch(Vec<Message>),
}

// Type aliases for cleaner usage
pub type MessageBroker = Broker<Message>;
pub type TopicName = String;
pub type ClientId = String;

pub struct Client {
    pub id: ClientId,
    pub subscribed_topics: Vec<TopicName>,
}

pub struct Broker<T> {
    pub topics: HashMap<String, VecDeque<T>>,
    pub capacity: usize,
}

pub fn publish<T>(broker: &mut Broker<T>, topic: &str, msg: T) {
    let q = broker.topics
        .entry(topic.to_string())
        .or_default();
        
    if q.len() >= broker.capacity {
        q.pop_front();
    }
    q.push_back(msg);
}

pub fn next_message<T>(broker: &mut Broker<T>, topic: &str) -> Option<T> {
    broker.topics.get_mut(topic)?.pop_front()
}

pub fn peek<T: Clone>(broker: &Broker<T>, topic: &str) -> Option<T> {
    broker.topics.get(topic)?.front().cloned()
}

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

    #[test]
    fn test_broker() {
        // Test custom message broker
        let mut broker = MessageBroker { topics: HashMap::new(), capacity: 2 };
        
        // Test text message
        publish(&mut broker, "chat", Message::Text("hello".to_string()));
        assert_eq!(
            next_message(&mut broker, "chat"), 
            Some(Message::Text("hello".to_string()))
        );
        
        // Test number and capacity limit
        publish(&mut broker, "metrics", Message::Number(100));
        publish(&mut broker, "metrics", Message::Number(200));
        publish(&mut broker, "metrics", Message::Number(300));
        assert_eq!(next_message(&mut broker, "metrics"), Some(Message::Number(200)));
        assert_eq!(next_message(&mut broker, "metrics"), Some(Message::Number(300)));
        
        // Test binary data with peek
        let binary_data = Message::Binary(vec![1, 2, 3]);
        publish(&mut broker, "data", binary_data.clone());
        assert_eq!(peek(&broker, "data"), Some(binary_data.clone()));
        
        // Test batch message
        let batch = Message::Batch(vec![
            Message::Text("first".to_string()),
            Message::Number(42),
        ]);
        publish(&mut broker, "batch", batch.clone());
        assert_eq!(next_message(&mut broker, "batch"), Some(batch));
    }

    // Original tests for basic types
    #[test]
    fn test_basic_types() {
        // Test integer messages
        let mut broker = Broker { topics: HashMap::new(), capacity: 2 };
        publish(&mut broker, "t1", 1);
        assert_eq!(next_message(&mut broker, "t1"), Some(1));
        
        // Test string messages with capacity limit
        let mut broker: Broker<String> = Broker { topics: HashMap::new(), capacity: 2 };
        publish(&mut broker, "t2", "a".to_string());
        publish(&mut broker, "t2", "b".to_string());
        publish(&mut broker, "t2", "c".to_string());
        assert_eq!(next_message(&mut broker, "t2"), Some("b".to_string()));
        assert_eq!(next_message(&mut broker, "t2"), Some("c".to_string()));
        assert_eq!(next_message(&mut broker, "t2"), None);
        
        // Test vec messages with peek
        let mut broker: Broker<Vec<i32>> = Broker { topics: HashMap::new(), capacity: 2 };
        publish(&mut broker, "t3", vec![1,2,3]);
        assert_eq!(peek(&broker, "t3"), Some(vec![1,2,3]));
        assert_eq!(next_message(&mut broker, "t3"), Some(vec![1,2,3]));
        assert_eq!(peek(&broker, "t3"), None);
    }
}