zulip_bots: Fix and simplify extract_query_without_mention.

This fixes quirks related to the regex not covering all potential
weird usernames and adds tests for stripping the @-mentions.
This commit is contained in:
derAnfaenger 2017-11-11 00:58:57 +01:00 committed by showell
parent 0a70a6b46f
commit 55332d8cbc
2 changed files with 26 additions and 6 deletions

View file

@ -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

View file

@ -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()