Skip to main content

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::TokioRuntime;
27    use crate::types::message::MessageInfo;
28    use anyhow::Result;
29    use async_lock::Mutex;
30    use std::sync::Arc;
31    use wacore_binary::node::Node;
32
33    /// Mock handler for testing custom enc types
34    #[derive(Debug)]
35    struct MockEncHandler {
36        pub calls: Arc<Mutex<Vec<String>>>,
37    }
38
39    impl MockEncHandler {
40        fn new() -> Self {
41            Self {
42                calls: Arc::new(Mutex::new(Vec::new())),
43            }
44        }
45    }
46
47    #[async_trait::async_trait]
48    impl EncHandler for MockEncHandler {
49        async fn handle(
50            &self,
51            _client: Arc<crate::client::Client>,
52            enc_node: &Node,
53            _info: &MessageInfo,
54        ) -> Result<()> {
55            let enc_type = enc_node
56                .attrs()
57                .optional_string("type")
58                .as_deref()
59                .unwrap_or("unknown")
60                .to_string();
61            self.calls.lock().await.push(enc_type);
62            Ok(())
63        }
64    }
65
66    #[tokio::test]
67    async fn test_custom_enc_handler_registration() {
68        use crate::bot::Bot;
69
70        // Create a mock handler
71        let mock_handler = MockEncHandler::new();
72
73        // Build bot with custom handler and in-memory DB
74        let backend = crate::test_utils::create_test_backend().await;
75
76        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
77        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
78        let bot = Bot::builder()
79            .with_backend(backend)
80            .with_transport_factory(transport)
81            .with_http_client(http_client)
82            .with_enc_handler("frskmsg", mock_handler)
83            .with_runtime(TokioRuntime)
84            .build()
85            .await
86            .expect("Failed to build bot");
87
88        // Verify handler was registered
89        assert!(
90            bot.client()
91                .custom_enc_handlers
92                .read()
93                .await
94                .contains_key("frskmsg")
95        );
96    }
97
98    #[tokio::test]
99    async fn test_multiple_custom_handlers() {
100        use crate::bot::Bot;
101
102        let handler1 = MockEncHandler::new();
103        let handler2 = MockEncHandler::new();
104
105        // Build bot with in-memory DB
106        let backend = crate::test_utils::create_test_backend().await;
107
108        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
109        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
110        let bot = Bot::builder()
111            .with_backend(backend)
112            .with_transport_factory(transport)
113            .with_http_client(http_client)
114            .with_enc_handler("frskmsg", handler1)
115            .with_enc_handler("customtype", handler2)
116            .with_runtime(TokioRuntime)
117            .build()
118            .await
119            .expect("Failed to build bot");
120
121        // Verify both handlers were registered
122        assert!(
123            bot.client()
124                .custom_enc_handlers
125                .read()
126                .await
127                .contains_key("frskmsg")
128        );
129        assert!(
130            bot.client()
131                .custom_enc_handlers
132                .read()
133                .await
134                .contains_key("customtype")
135        );
136        assert_eq!(bot.client().custom_enc_handlers.read().await.len(), 2);
137    }
138
139    #[tokio::test]
140    async fn test_builtin_handlers_still_work() {
141        use crate::bot::Bot;
142
143        // Build bot without custom handlers but with in-memory DB
144        let backend = crate::test_utils::create_test_backend().await;
145
146        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
147        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
148        let bot = Bot::builder()
149            .with_backend(backend)
150            .with_transport_factory(transport)
151            .with_http_client(http_client)
152            .with_runtime(TokioRuntime)
153            .build()
154            .await
155            .expect("Failed to build bot");
156
157        // Verify no custom handlers are registered
158        assert_eq!(bot.client().custom_enc_handlers.read().await.len(), 0);
159    }
160}