kernel_sidecar/handlers/
msg_count.rs

1use std::collections::HashMap;
2
3use crate::jupyter::response::Response;
4
5use crate::handlers::Handler;
6
7// Returns a hashmap of {msg_type: count} for all messages handled by an Action
8// Primarily used in tests and introspective click-testing
9#[derive(Debug, Clone)]
10pub struct MessageCountHandler {
11    pub counts: HashMap<String, usize>,
12}
13
14impl Default for MessageCountHandler {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl MessageCountHandler {
21    pub fn new() -> Self {
22        MessageCountHandler {
23            counts: HashMap::new(),
24        }
25    }
26}
27
28#[async_trait::async_trait]
29impl Handler for MessageCountHandler {
30    async fn handle(&mut self, msg: &Response) {
31        let msg_type = msg.msg_type();
32        let count = self.counts.entry(msg_type).or_insert(0);
33        *count += 1;
34    }
35}