use chrono::{TimeZone, Utc}; use mxlink::matrix_sdk::ruma::OwnedUserId; use crate::conversation::matrix::{ MatrixMessage, MatrixMessageContent, MatrixMessageProcessingParams, }; #[test] fn is_message_from_allowed_sender() { let bot_user_id = OwnedUserId::try_from("@bot:example.com").expect("Failed to parse bot user ID"); let allowed_user_id = OwnedUserId::try_from("@user.someone:example.com").unwrap(); let unallowed_user_id = OwnedUserId::try_from("@another:example.com").unwrap(); let timestamp = Utc.with_ymd_and_hms(2024, 9, 20, 18, 34, 15).unwrap(); let bot_message = MatrixMessage { sender_id: bot_user_id.to_owned(), content: MatrixMessageContent::Text("Hello!".to_owned()), mentioned_users: vec![], timestamp, }; let allowed_user_message = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: MatrixMessageContent::Text("Hello!".to_owned()), mentioned_users: vec![], timestamp, }; let unallowed_user_message = MatrixMessage { sender_id: unallowed_user_id.to_owned(), content: MatrixMessageContent::Text("Hello!".to_owned()), mentioned_users: vec![], timestamp, }; let parsed_regex = match mxidwc::parse_pattern("@user.*:example.com") { Ok(value) => value, Err(err) => { panic!("Error parsing regex: {}", err); } }; let allowed_users = vec![parsed_regex]; assert!( super::is_message_from_allowed_sender(&bot_message, &bot_user_id, Some(&allowed_users)), "Bot message should be allowed" ); assert!( super::is_message_from_allowed_sender( &allowed_user_message, &bot_user_id, Some(&allowed_users) ), "Allowed user message should be allowed" ); assert!( !super::is_message_from_allowed_sender( &unallowed_user_message, &bot_user_id, Some(&allowed_users), ), "Unallowed user message should be ignored" ); assert!( super::is_message_from_allowed_sender(&unallowed_user_message, &bot_user_id, None,), "An empty list of allowed users lets everyone through" ); } #[tokio::test] async fn process_matrix_messages() { let bot_user_id = OwnedUserId::try_from("@bot:example.com").expect("Failed to parse bot user ID"); let allowed_user_id = OwnedUserId::try_from("@user.someone:example.com").unwrap(); let unallowed_user_id = OwnedUserId::try_from("@another:example.com").unwrap(); let timestamp = Utc.with_ymd_and_hms(2024, 9, 20, 18, 34, 15).unwrap(); let allowed_user_message = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: MatrixMessageContent::Text("Hello from the user!".to_owned()), mentioned_users: vec![], timestamp, }; let allowed_user_message_with_prefix = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: MatrixMessageContent::Text("!bai Hello from the user!".to_owned()), mentioned_users: vec![], timestamp, }; let allowed_user_message_with_prefix_no_space = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: MatrixMessageContent::Text("!baiHello from the user!".to_owned()), mentioned_users: vec![], timestamp, }; let allowed_user_message_with_prefix_full_width_space = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: MatrixMessageContent::Text("!bai Hello from the user!".to_owned()), mentioned_users: vec![], timestamp, }; let bot_message = MatrixMessage { sender_id: bot_user_id.to_owned(), content: MatrixMessageContent::Text("Hello from the bot!".to_owned()), mentioned_users: vec![], timestamp, }; let allowed_user_message_with_bot_mention = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: MatrixMessageContent::Text("@baibot: Hello from the user!".to_owned()), mentioned_users: vec![bot_user_id.to_owned()], timestamp, }; // The message text is the same as above - it mentions the bot, but the actually-mentioned user is another user. let allowed_user_message_with_another_user_mention = MatrixMessage { sender_id: allowed_user_id.to_owned(), content: allowed_user_message_with_bot_mention.content.clone(), mentioned_users: vec![allowed_user_id.to_owned()], timestamp, }; let unallowed_user_message = MatrixMessage { sender_id: unallowed_user_id.to_owned(), content: MatrixMessageContent::Text("Hello from an unallowed user!".to_owned()), mentioned_users: vec![], timestamp, }; let parsed_regex = match mxidwc::parse_pattern("@user.*:example.com") { Ok(value) => value, Err(err) => { panic!("Error parsing regex: {}", err); } }; let allowed_users = vec![parsed_regex]; let message_processing_params_basic = super::MatrixMessageProcessingParams::new( bot_user_id.to_owned(), Some(allowed_users.clone()), ); let message_processing_params_with_prefix_stripping = super::MatrixMessageProcessingParams::new( bot_user_id.to_owned(), Some(allowed_users.clone()), ) .with_first_message_prefixes_to_strip(vec!["!bai".to_owned()]); let message_processing_params_with_bot_user_prefix_stripping = super::MatrixMessageProcessingParams::new( bot_user_id.to_owned(), Some(allowed_users.clone()), ) .with_bot_user_prefixes_to_strip(vec!["@baibot: ".to_owned(), "@baibot".to_owned()]); struct TestCase { name: String, messages: Vec, message_processing_params: MatrixMessageProcessingParams, expected_message_texts: Vec, } let test_cases = vec![ TestCase { name: "Messages by unallowed users are ignored".to_owned(), messages: vec![ allowed_user_message.clone(), bot_message.clone(), unallowed_user_message.clone(), ], message_processing_params: message_processing_params_basic.clone(), expected_message_texts: vec![ "Hello from the user!".to_owned(), "Hello from the bot!".to_owned(), ], }, TestCase { name: "The first message with a prefix gets stripped if params configure it (regular space)".to_owned(), messages: vec![ allowed_user_message_with_prefix.clone(), bot_message.clone(), allowed_user_message_with_prefix.clone(), unallowed_user_message.clone(), ], message_processing_params: message_processing_params_with_prefix_stripping.clone(), expected_message_texts: vec![ "Hello from the user!".to_owned(), "Hello from the bot!".to_owned(), "!bai Hello from the user!".to_owned(), ], }, TestCase { name: "The first message with a prefix gets stripped if params configure it (no space)".to_owned(), messages: vec![ allowed_user_message_with_prefix_no_space.clone(), bot_message.clone(), allowed_user_message_with_prefix_no_space.clone(), unallowed_user_message.clone(), ], message_processing_params: message_processing_params_with_prefix_stripping.clone(), expected_message_texts: vec![ "Hello from the user!".to_owned(), "Hello from the bot!".to_owned(), "!baiHello from the user!".to_owned(), ], }, TestCase { name: "The first message with a prefix gets stripped if params configure it (full-width-space)".to_owned(), messages: vec![ allowed_user_message_with_prefix_full_width_space.clone(), bot_message.clone(), allowed_user_message_with_prefix_full_width_space.clone(), unallowed_user_message.clone(), ], message_processing_params: message_processing_params_with_prefix_stripping.clone(), expected_message_texts: vec![ "Hello from the user!".to_owned(), "Hello from the bot!".to_owned(), "!bai Hello from the user!".to_owned(), ], }, TestCase { name: "The first message with a prefix remains untouched if params leave it alone" .to_owned(), messages: vec![ allowed_user_message_with_prefix.clone(), bot_message.clone(), allowed_user_message_with_prefix.clone(), unallowed_user_message.clone(), ], message_processing_params: message_processing_params_basic.clone(), expected_message_texts: vec![ "!bai Hello from the user!".to_owned(), "Hello from the bot!".to_owned(), "!bai Hello from the user!".to_owned(), ], }, TestCase { name: "Messages that mention the bot user get the bot user prefix stripped" .to_owned(), messages: vec![ allowed_user_message_with_bot_mention.clone(), allowed_user_message_with_another_user_mention.clone(), ], message_processing_params: message_processing_params_with_bot_user_prefix_stripping.clone(), expected_message_texts: vec![ "Hello from the user!".to_owned(), "@baibot: Hello from the user!".to_owned(), ], }, ]; for test_case in test_cases { let processed_messages = super::process_matrix_messages( &test_case.messages, &test_case.message_processing_params, ) .await; let processed_message_texts = processed_messages .iter() .map(|message| match &message.content { MatrixMessageContent::Text(text) => text.clone(), _ => "".to_owned(), }) .collect::>(); assert_eq!( processed_message_texts, test_case.expected_message_texts, "Test case {} failed", test_case.name, ); } } #[test] fn create_list_of_bot_user_prefixes_to_strip() { let bot_user_id = OwnedUserId::try_from("@baibot:example.com").expect("Failed to parse bot user ID"); // Test case 1: Bot user with no display name let bot_display_name = None; let prefixes = super::create_list_of_bot_user_prefixes_to_strip(&bot_user_id, &bot_display_name); assert_eq!( prefixes, vec![ "@baibot:example.com".to_string(), "@baibot".to_string(), "baibot".to_string(), ":".to_string() ] ); // Test case 2: Bot user with display name let bot_display_name = Some("Assistant".to_string()); let prefixes = super::create_list_of_bot_user_prefixes_to_strip(&bot_user_id, &bot_display_name); assert_eq!( prefixes, vec![ "@baibot:example.com".to_string(), "@baibot".to_string(), "baibot".to_string(), "@Assistant".to_string(), "Assistant".to_string(), ":".to_string() ] ); }