From 55332d8cbc24c81fa16d3c89b8a2a879a06040c5 Mon Sep 17 00:00:00 2001
From: derAnfaenger <robhoenig@gmail.com>
Date: Sat, 11 Nov 2017 00:58:57 +0100
Subject: [PATCH] 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.
---
 zulip_bots/zulip_bots/lib.py      | 10 ++++------
 zulip_bots/zulip_bots/test_run.py | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 6 deletions(-)

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