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::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
55                .attrs()
56                .optional_string("type")
57                .unwrap_or("unknown")
58                .to_string();
59            self.calls.lock().await.push(enc_type);
60            Ok(())
61        }
62    }
63
64    #[tokio::test]
65    async fn test_custom_enc_handler_registration() {
66        use crate::bot::Bot;
67
68        // Create a mock handler
69        let mock_handler = MockEncHandler::new();
70
71        // Build bot with custom handler and in-memory DB
72        let backend = crate::test_utils::create_test_backend().await;
73
74        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
75        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
76        let bot = Bot::builder()
77            .with_backend(backend)
78            .with_transport_factory(transport)
79            .with_http_client(http_client)
80            .with_enc_handler("frskmsg", mock_handler)
81            .build()
82            .await
83            .expect("Failed to build bot");
84
85        // Verify handler was registered
86        assert!(bot.client().custom_enc_handlers.contains_key("frskmsg"));
87    }
88
89    #[tokio::test]
90    async fn test_multiple_custom_handlers() {
91        use crate::bot::Bot;
92
93        let handler1 = MockEncHandler::new();
94        let handler2 = MockEncHandler::new();
95
96        // Build bot with in-memory DB
97        let backend = crate::test_utils::create_test_backend().await;
98
99        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
100        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
101        let bot = Bot::builder()
102            .with_backend(backend)
103            .with_transport_factory(transport)
104            .with_http_client(http_client)
105            .with_enc_handler("frskmsg", handler1)
106            .with_enc_handler("customtype", handler2)
107            .build()
108            .await
109            .expect("Failed to build bot");
110
111        // Verify both handlers were registered
112        assert!(bot.client().custom_enc_handlers.contains_key("frskmsg"));
113        assert!(bot.client().custom_enc_handlers.contains_key("customtype"));
114        assert_eq!(bot.client().custom_enc_handlers.len(), 2);
115    }
116
117    #[tokio::test]
118    async fn test_builtin_handlers_still_work() {
119        use crate::bot::Bot;
120
121        // Build bot without custom handlers but with in-memory DB
122        let backend = crate::test_utils::create_test_backend().await;
123
124        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
125        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
126        let bot = Bot::builder()
127            .with_backend(backend)
128            .with_transport_factory(transport)
129            .with_http_client(http_client)
130            .build()
131            .await
132            .expect("Failed to build bot");
133
134        // Verify no custom handlers are registered
135        assert_eq!(bot.client().custom_enc_handlers.len(), 0);
136    }
137}