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::SqliteStore::new(":memory:")
70                .await
71                .expect("Failed to create SQLite backend"),
72        );
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 = Arc::new(
98            crate::store::SqliteStore::new(":memory:")
99                .await
100                .expect("Failed to create SQLite backend"),
101        );
102
103        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
104        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
105        let bot = Bot::builder()
106            .with_backend(backend)
107            .with_transport_factory(transport)
108            .with_http_client(http_client)
109            .with_enc_handler("frskmsg", handler1)
110            .with_enc_handler("customtype", handler2)
111            .build()
112            .await
113            .expect("Failed to build bot");
114
115        // Verify both handlers were registered
116        assert!(bot.client().custom_enc_handlers.contains_key("frskmsg"));
117        assert!(bot.client().custom_enc_handlers.contains_key("customtype"));
118        assert_eq!(bot.client().custom_enc_handlers.len(), 2);
119    }
120
121    #[tokio::test]
122    async fn test_builtin_handlers_still_work() {
123        use crate::bot::Bot;
124
125        // Build bot without custom handlers but with in-memory DB
126        let backend = Arc::new(
127            crate::store::SqliteStore::new(":memory:")
128                .await
129                .expect("Failed to create SQLite backend"),
130        );
131
132        let transport = whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory::new();
133        let http_client = whatsapp_rust_ureq_http_client::UreqHttpClient::new();
134        let bot = Bot::builder()
135            .with_backend(backend)
136            .with_transport_factory(transport)
137            .with_http_client(http_client)
138            .build()
139            .await
140            .expect("Failed to build bot");
141
142        // Verify no custom handlers are registered
143        assert_eq!(bot.client().custom_enc_handlers.len(), 0);
144    }
145}
whatsapp_rust/types/enc_handler.rs

whatsapp_rust/types/
enc_handler.rs
