diff --git a/zulip_bots/zulip_bots/lib.py b/zulip_bots/zulip_bots/lib.py index f4a4b08..d934aca 100644 --- a/zulip_bots/zulip_bots/lib.py +++ b/zulip_bots/zulip_bots/lib.py @@ -5,6 +5,7 @@ import os import signal import sys import time +import re from typing import Any, Optional, List, Dict, IO, Text @@ -238,10 +239,18 @@ def extract_query_without_mention(message: Dict[str, Any], client: ExternalBotHa If the bot is the first @mention in the message, then this function returns the stripped message with the bot's @mention removed. Otherwise, it returns None. """ + content = message['content'] mention = '@**' + client.full_name + '**' - if not message['content'].startswith(mention): - return None - return message['content'][len(mention):].lstrip() + extended_mention_regex = re.compile(r'^@\*\*.*\|' + str(client.user_id) + r'\*\*') + extended_mention_match = extended_mention_regex.match(content) + + if extended_mention_match: + return content[extended_mention_match.end():].lstrip() + + if content.startswith(mention): + return content[len(mention):].lstrip() + + return None def is_private_message_from_another_user(message_dict: Dict[str, Any], current_user_id: int) -> bool: diff --git a/zulip_bots/zulip_bots/tests/test_lib.py b/zulip_bots/zulip_bots/tests/test_lib.py index 5c0ac30..dbf882e 100644 --- a/zulip_bots/zulip_bots/tests/test_lib.py +++ b/zulip_bots/zulip_bots/tests/test_lib.py @@ -4,6 +4,7 @@ from zulip_bots.lib import ( ExternalBotHandler, StateHandler, run_message_handler_for_bot, + extract_query_without_mention ) import io @@ -192,6 +193,23 @@ class LibTest(TestCase): client.upload_file.assert_called_once_with(file) + def test_extract_query_without_mention(self): + client = FakeClient() + handler = ExternalBotHandler( + client=client, + root_dir=None, + bot_details=None, + bot_config_file=None + ) + message = {'content': "@**Alice** Hello World"} + self.assertEqual(extract_query_without_mention(message, handler), "Hello World") + message = {'content': "@**Alice|alice** Hello World"} + self.assertEqual(extract_query_without_mention(message, handler), "Hello World") + message = {'content': "@**Alice Renamed|alice** Hello World"} + self.assertEqual(extract_query_without_mention(message, handler), "Hello World") + message = {'content': "Not at start @**Alice|alice** Hello World"} + self.assertEqual(extract_query_without_mention(message, handler), None) + def _create_client_and_handler_for_file_upload(self): client = FakeClient() client.upload_file = MagicMock()