whatsapp_rust/types/
enc_handler.rs

1use crate::client::Client;
2use crate::types::message::MessageInfo;
3use anyhow::Result;
4use std::sync::Arc;
5use wacore_binary::node::Node;
6
7/// Trait for handling custom encrypted message types
8#[async_trait::async_trait]
9pub trait EncHandler: Send + Sync {
10    /// Handle an encrypted node of a specific type
11    ///
12    /// # Arguments
13    /// * `client` - The client instance
14    /// * `enc_node` - The encrypted node to handle
15    /// * `info` - The message info context
16    ///
17    /// # Returns
18    /// * `Ok(())` if the message was handled successfully
19    /// * `Err(anyhow::Error)` if handling failed
20    async fn handle(&self, client: Arc<Client>, enc_node: &Node, info: &MessageInfo) -> Result<()>;
21}
22
23#[cfg(test)]
24mod tests {
25    use super::*;
26    use crate::types::message::MessageInfo;
27    use anyhow::Result;
28    use std::sync::Arc;
29    use tokio::sync::Mutex;
30    use wacore_binary::node::Node;
31
32    /// Mock handler for testing custom enc types
33    #[derive(Debug)]
34    struct MockEncHandler {
35        pub calls: Arc<Mutex<Vec<String>>>,
36    }
37
38    impl MockEncHandler {
39        fn new() -> Self {
40            Self {
41                calls: Arc::new(Mutex::new(Vec::new())),
42            }
43        }
44    }
45
46    #[async_trait::async_trait]
47    impl EncHandler for MockEncHandler {
48        async fn handle(
49            &self,
50            _client: Arc<crate::client::Client>,
51            enc_node: &Node,
52            _info: &MessageInfo,
53        ) -> Result<()> {
54            let enc_type = enc_node.attrs().string("type");
55            self.calls.lock().await.push(enc_type);
56            Ok(())
57        }
58    }
59
60    #[tokio::test]
61    async fn test_custom_enc_handler_registration() {
62        use crate::bot::Bot;
63
64        // Create a mock handler
65        let mock_handler = MockEncHandler::new();
66
67        // Build bot with custom handler and in-memory DB
68        let backend = Arc::new(
69            crate::store::sqlite_store::SqliteStore::new(":memory:")
70                .await
71                .expect("Failed to create SQLite backend"),
72        );
73
74        let bot = Bot::builder()
75            .with_backend(backend)
76            .with_enc_handler("frskmsg", mock_handler)
77            .build()
78            .await
79            .expect("Failed to build bot");
80
81        // Verify handler was registered
82        assert!(bot.client().custom_enc_handlers.contains_key("frskmsg"));
83    }
84
85    #[tokio::test]
86    async fn test_multiple_custom_handlers() {
87        use crate::bot::Bot;
88
89        let handler1 = MockEncHandler::new();
90        let handler2 = MockEncHandler::new();
91
92        // Build bot with in-memory DB
93        let backend = Arc::new(
94            crate::store::sqlite_store::SqliteStore::new(":memory:")
95                .await
96                .expect("Failed to create SQLite backend"),
97        );
98
99        let bot = Bot::builder()
100            .with_backend(backend)
101            .with_enc_handler("frskmsg", handler1)
102            .with_enc_handler("customtype", handler2)
103            .build()
104            .await
105            .expect("Failed to build bot");
106
107        // Verify both handlers were registered
108        assert!(bot.client().custom_enc_handlers.contains_key("frskmsg"));
109        assert!(bot.client().custom_enc_handlers.contains_key("customtype"));
110        assert_eq!(bot.client().custom_enc_handlers.len(), 2);
111    }
112
113    #[tokio::test]
114    async fn test_builtin_handlers_still_work() {
115        use crate::bot::Bot;
116
117        // Build bot without custom handlers but with in-memory DB
118        let backend = Arc::new(
119            crate::store::sqlite_store::SqliteStore::new(":memory:")
120                .await
121                .expect("Failed to create SQLite backend"),
122        );
123
124        let bot = Bot::builder()
125            .with_backend(backend)
126            .build()
127            .await
128            .expect("Failed to build bot");
129
130        // Verify no custom handlers are registered
131        assert_eq!(bot.client().custom_enc_handlers.len(), 0);
132    }
133}