diff --git a/zulip_bots/zulip_bots/lib.py b/zulip_bots/zulip_bots/lib.py index 4f9daaf..d8cb0bc 100644 --- a/zulip_bots/zulip_bots/lib.py +++ b/zulip_bots/zulip_bots/lib.py @@ -146,14 +146,12 @@ def extract_query_without_mention(message, client): # type: (Dict[str, Any], ExternalBotHandler) -> str """ If the bot is the first @mention in the message, then this function returns - the message with the bot's @mention removed. Otherwise, it returns None. + the stripped message with the bot's @mention removed. Otherwise, it returns None. """ - bot_mention = r'^@(\*\*{0}\*\*)'.format(client.full_name) - start_with_mention = re.compile(bot_mention).match(message['content']) - if start_with_mention is None: + mention = '@**' + client.full_name + '**' + if not message['content'].startswith(mention): return None - query_without_mention = message['content'][len(start_with_mention.group()):] - return query_without_mention.lstrip() + return message['content'][len(mention):].lstrip() def is_private_message_from_another_user(message_dict, current_user_id): # type: (Dict[str, Any], int) -> bool diff --git a/zulip_bots/zulip_bots/test_run.py b/zulip_bots/zulip_bots/test_run.py index 3984934..4b4ed63 100644 --- a/zulip_bots/zulip_bots/test_run.py +++ b/zulip_bots/zulip_bots/test_run.py @@ -4,11 +4,13 @@ from __future__ import absolute_import import importlib import os import zulip_bots.run +from zulip_bots.lib import extract_query_without_mention import six import unittest import zulip from importlib import import_module +from typing import Optional from unittest import TestCase if six.PY2: @@ -44,5 +46,25 @@ class TestDefaultArguments(TestCase): lib_module=mock.ANY, quiet=False) +class TestBotLib(TestCase): + def test_extract_query_without_mention(self): + # type: () -> None + + def test_message(name, message, expected_return): + # type: (str, str, Optional[str]) -> None + mock_client = mock.MagicMock() + mock_client.full_name = name + mock_message = {'content': message} + self.assertEqual(expected_return, extract_query_without_mention(mock_message, mock_client)) + test_message("xkcd", "@**xkcd**foo", "foo") + test_message("xkcd", "@**xkcd** foo", "foo") + test_message("xkcd", "@**xkcd** foo bar baz", "foo bar baz") + test_message("xkcd", "@**xkcd** foo bar baz", "foo bar baz") + test_message("xkcd", "@**xkcd** 123_) (/&%) +}}}l", "123_) (/&%) +}}}l") + test_message("brokenmention", "@**brokenmention* foo", None) + test_message("nomention", "foo", None) + test_message("Max Mustermann", "@**Max Mustermann** foo", "foo") + test_message("Max (Mustermann)#(*$&12]\]", "@**Max (Mustermann)#(*$&12]\]** foo", "foo") + if __name__ == '__main__': unittest.main()