zulip/zerver/tests/test_service_bot_system.py

# -*- coding: utf-8 -*-

import json
import mock
from typing import Any, Union, Mapping, Callable

from django.conf import settings
from django.test import override_settings

from zerver.lib.actions import (
    do_create_user,
    get_service_bot_events,
)
from zerver.lib.bot_lib import StateHandler, EmbeddedBotHandler
from zerver.lib.bot_storage import StateError
from zerver.lib.bot_config import set_bot_config, ConfigError
from zerver.lib.test_classes import ZulipTestCase
from zerver.models import (
    get_realm,
    BotUserStateData,
    UserProfile,
    Recipient,
)

BOT_TYPE_TO_QUEUE_NAME = {
    UserProfile.OUTGOING_WEBHOOK_BOT: 'outgoing_webhooks',
    UserProfile.EMBEDDED_BOT: 'embedded_bots',
}

class TestServiceBotBasics(ZulipTestCase):
    def _get_outgoing_bot(self) -> UserProfile:
        outgoing_bot = do_create_user(
            email="bar-bot@zulip.com",
            password="test",
            realm=get_realm("zulip"),
            full_name="BarBot",
            short_name='bb',
            bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
            bot_owner=self.example_user('cordelia'),
        )

        return outgoing_bot

    def test_service_events_for_pms(self) -> None:
        sender = self.example_user('hamlet')
        assert(not sender.is_bot)

        outgoing_bot = self._get_outgoing_bot()

        event_dict = get_service_bot_events(
            sender=sender,
            service_bot_tuples=[
                (outgoing_bot.id, outgoing_bot.bot_type),
            ],
            active_user_ids={outgoing_bot.id},
            mentioned_user_ids=set(),
            recipient_type=Recipient.PERSONAL,
        )

        expected = dict(
            outgoing_webhooks=[
                dict(trigger='private_message', user_profile_id=outgoing_bot.id),
            ],
        )

        self.assertEqual(event_dict, expected)

    def test_service_events_for_stream_mentions(self) -> None:
        sender = self.example_user('hamlet')
        assert(not sender.is_bot)

        outgoing_bot = self._get_outgoing_bot()

        event_dict = get_service_bot_events(
            sender=sender,
            service_bot_tuples=[
                (outgoing_bot.id, outgoing_bot.bot_type),
            ],
            active_user_ids=set(),
            mentioned_user_ids={outgoing_bot.id},
            recipient_type=Recipient.STREAM,
        )

        expected = dict(
            outgoing_webhooks=[
                dict(trigger='mention', user_profile_id=outgoing_bot.id),
            ],
        )

        self.assertEqual(event_dict, expected)

class TestServiceBotStateHandler(ZulipTestCase):
    def setUp(self) -> None:
        self.user_profile = self.example_user("othello")
        self.bot_profile = do_create_user(email="embedded-bot-1@zulip.com",
                                          password="test",
                                          realm=get_realm("zulip"),
                                          full_name="EmbeddedBo1",
                                          short_name="embedded-bot-1",
                                          bot_type=UserProfile.EMBEDDED_BOT,
                                          bot_owner=self.user_profile)
        self.second_bot_profile = do_create_user(email="embedded-bot-2@zulip.com",
                                                 password="test",
                                                 realm=get_realm("zulip"),
                                                 full_name="EmbeddedBot2",
                                                 short_name="embedded-bot-2",
                                                 bot_type=UserProfile.EMBEDDED_BOT,
                                                 bot_owner=self.user_profile)

    def test_basic_storage_and_retrieval(self) -> None:
        storage = StateHandler(self.bot_profile)
        storage.put('some key', 'some value')
        storage.put('some other key', 'some other value')
        self.assertEqual(storage.get('some key'), 'some value')
        self.assertEqual(storage.get('some other key'), 'some other value')
        self.assertTrue(storage.contains('some key'))
        self.assertFalse(storage.contains('nonexistent key'))
        self.assertRaisesMessage(StateError,
                                 "Cannot get state. <UserProfile: embedded-bot-1@zulip.com <Realm: zulip 1>> "
                                 "doesn't have an entry with the key 'nonexistent key'.",
                                 lambda: storage.get('nonexistent key'))
        storage.put('some key', 'a new value')
        self.assertEqual(storage.get('some key'), 'a new value')
        second_storage = StateHandler(self.second_bot_profile)
        self.assertRaises(StateError, lambda: second_storage.get('some key'))
        second_storage.put('some key', 'yet another value')
        self.assertEqual(storage.get('some key'), 'a new value')
        self.assertEqual(second_storage.get('some key'), 'yet another value')

    def test_marshaling(self) -> None:
        storage = StateHandler(self.bot_profile)
        serializable_obj = {'foo': 'bar', 'baz': [42, 'cux']}
        storage.put('some key', serializable_obj)  # type: ignore # Ignore for testing.
        self.assertEqual(storage.get('some key'), serializable_obj)

    def test_invalid_calls(self) -> None:
        storage = StateHandler(self.bot_profile)
        storage.marshal = lambda obj: obj
        storage.demarshal = lambda obj: obj
        serializable_obj = {'foo': 'bar', 'baz': [42, 'cux']}
        with self.assertRaisesMessage(StateError, "Cannot set state. The value type is "
                                                  "<class 'dict'>, but it should be str."):
            storage.put('some key', serializable_obj)  # type: ignore # We intend to test an invalid type.
        with self.assertRaisesMessage(StateError, "Cannot set state. The key type is "
                                                  "<class 'dict'>, but it should be str."):
            storage.put(serializable_obj, 'some value')  # type: ignore # We intend to test an invalid type.

    # Reduce maximal state size for faster test string construction.
    @override_settings(USER_STATE_SIZE_LIMIT=100)
    def test_storage_limit(self) -> None:
        storage = StateHandler(self.bot_profile)

        # Disable marshaling for storing a string whose size is
        # equivalent to the size of the stored object.
        storage.marshal = lambda obj: obj
        storage.demarshal = lambda obj: obj

        key = 'capacity-filling entry'
        storage.put(key, 'x' * (settings.USER_STATE_SIZE_LIMIT - len(key)))

        with self.assertRaisesMessage(StateError, "Cannot set state. Request would require 132 bytes storage. "
                                                  "The current storage limit is 100."):
            storage.put('too much data', 'a few bits too long')

        second_storage = StateHandler(self.second_bot_profile)
        second_storage.put('another big entry', 'x' * (settings.USER_STATE_SIZE_LIMIT - 40))
        second_storage.put('normal entry', 'abcd')

    def test_entry_removal(self) -> None:
        storage = StateHandler(self.bot_profile)
        storage.put('some key', 'some value')
        storage.put('another key', 'some value')
        self.assertTrue(storage.contains('some key'))
        self.assertTrue(storage.contains('another key'))
        storage.remove('some key')
        self.assertFalse(storage.contains('some key'))
        self.assertTrue(storage.contains('another key'))
        self.assertRaises(StateError, lambda: storage.remove('some key'))

class TestServiceBotConfigHandler(ZulipTestCase):
    def setUp(self) -> None:
        self.user_profile = self.example_user("othello")
        self.bot_profile = self.create_test_bot('embedded-bot@zulip.testserver', self.user_profile, 'Embedded bot',
                                                'embedded', UserProfile.EMBEDDED_BOT, service_name='helloworld')
        self.bot_handler = EmbeddedBotHandler(self.bot_profile)

    def test_basic_storage_and_retrieval(self) -> None:
        config_dict = {"entry 1": "value 1", "entry 2": "value 2"}
        for key, value in config_dict.items():
            set_bot_config(self.bot_profile, key, value)
        self.assertEqual(self.bot_handler.get_config_info(), config_dict)

        config_update = {"entry 2": "new value", "entry 3": "value 3"}
        for key, value in config_update.items():
            set_bot_config(self.bot_profile, key, value)
        config_dict.update(config_update)
        self.assertEqual(self.bot_handler.get_config_info(), config_dict)

    @override_settings(BOT_CONFIG_SIZE_LIMIT=100)
    def test_config_entry_limit(self) -> None:
        set_bot_config(self.bot_profile, "some key", 'x' * (settings.BOT_CONFIG_SIZE_LIMIT-8))
        self.assertRaisesMessage(ConfigError,
                                 "Cannot store configuration. Request would require 101 characters. "
                                 "The current configuration size limit is 100 characters.",
                                 lambda: set_bot_config(self.bot_profile, "some key", 'x' * (settings.BOT_CONFIG_SIZE_LIMIT-8+1)))
        set_bot_config(self.bot_profile, "some key", 'x' * (settings.BOT_CONFIG_SIZE_LIMIT-20))
        set_bot_config(self.bot_profile, "another key", 'x')
        self.assertRaisesMessage(ConfigError,
                                 "Cannot store configuration. Request would require 116 characters. "
                                 "The current configuration size limit is 100 characters.",
                                 lambda: set_bot_config(self.bot_profile, "yet another key", 'x'))

class TestServiceBotEventTriggers(ZulipTestCase):

    def setUp(self) -> None:
        self.user_profile = self.example_user("othello")
        self.bot_profile = do_create_user(email="foo-bot@zulip.com",
                                          password="test",
                                          realm=get_realm("zulip"),
                                          full_name="FooBot",
                                          short_name="foo-bot",
                                          bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
                                          bot_owner=self.user_profile)
        self.second_bot_profile = do_create_user(email="bar-bot@zulip.com",
                                                 password="test",
                                                 realm=get_realm("zulip"),
                                                 full_name="BarBot",
                                                 short_name="bar-bot",
                                                 bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
                                                 bot_owner=self.user_profile)

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_trigger_on_stream_mention_from_user(self, mock_queue_json_publish: mock.Mock) -> None:
        for bot_type, expected_queue_name in BOT_TYPE_TO_QUEUE_NAME.items():
            self.bot_profile.bot_type = bot_type
            self.bot_profile.save()

            content = u'@**FooBot** foo bar!!!'
            recipient = 'Denmark'
            trigger = 'mention'
            message_type = Recipient._type_names[Recipient.STREAM]

            def check_values_passed(queue_name: Any,
                                    trigger_event: Union[Mapping[Any, Any], Any],
                                    x: Callable[[Any], None],
                                    call_consume_in_tests: bool) -> None:
                self.assertEqual(queue_name, expected_queue_name)
                self.assertEqual(trigger_event["message"]["content"], content)
                self.assertEqual(trigger_event["message"]["display_recipient"], recipient)
                self.assertEqual(trigger_event["message"]["sender_email"], self.user_profile.email)
                self.assertEqual(trigger_event["message"]["type"], message_type)
                self.assertEqual(trigger_event['trigger'], trigger)
                self.assertEqual(trigger_event['user_profile_id'], self.bot_profile.id)
            mock_queue_json_publish.side_effect = check_values_passed

            self.send_stream_message(
                self.user_profile.email,
                'Denmark',
                content)
            self.assertTrue(mock_queue_json_publish.called)

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_no_trigger_on_stream_message_without_mention(self, mock_queue_json_publish: mock.Mock) -> None:
        sender_email = self.user_profile.email
        self.send_stream_message(sender_email, "Denmark")
        self.assertFalse(mock_queue_json_publish.called)

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_no_trigger_on_stream_mention_from_bot(self, mock_queue_json_publish: mock.Mock) -> None:
        for bot_type in BOT_TYPE_TO_QUEUE_NAME:
            self.bot_profile.bot_type = bot_type
            self.bot_profile.save()

            self.send_stream_message(
                self.second_bot_profile.email,
                'Denmark',
                u'@**FooBot** foo bar!!!')
            self.assertFalse(mock_queue_json_publish.called)

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_trigger_on_personal_message_from_user(self, mock_queue_json_publish: mock.Mock) -> None:
        for bot_type, expected_queue_name in BOT_TYPE_TO_QUEUE_NAME.items():
            self.bot_profile.bot_type = bot_type
            self.bot_profile.save()

            sender_email = self.user_profile.email
            recipient_email = self.bot_profile.email

            def check_values_passed(queue_name: Any,
                                    trigger_event: Union[Mapping[Any, Any], Any],
                                    x: Callable[[Any], None],
                                    call_consume_in_tests: bool) -> None:
                self.assertEqual(queue_name, expected_queue_name)
                self.assertEqual(trigger_event["user_profile_id"], self.bot_profile.id)
                self.assertEqual(trigger_event["trigger"], "private_message")
                self.assertEqual(trigger_event["message"]["sender_email"], sender_email)
                display_recipients = [
                    trigger_event["message"]["display_recipient"][0]["email"],
                    trigger_event["message"]["display_recipient"][1]["email"],
                ]
                self.assertTrue(sender_email in display_recipients)
                self.assertTrue(recipient_email in display_recipients)
            mock_queue_json_publish.side_effect = check_values_passed

            self.send_personal_message(sender_email, recipient_email, 'test')
            self.assertTrue(mock_queue_json_publish.called)

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_no_trigger_on_personal_message_from_bot(self, mock_queue_json_publish: mock.Mock) -> None:
        for bot_type in BOT_TYPE_TO_QUEUE_NAME:
            self.bot_profile.bot_type = bot_type
            self.bot_profile.save()

            sender_email = self.second_bot_profile.email
            recipient_email = self.bot_profile.email
            self.send_personal_message(sender_email, recipient_email)
            self.assertFalse(mock_queue_json_publish.called)

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_trigger_on_huddle_message_from_user(self, mock_queue_json_publish: mock.Mock) -> None:
        for bot_type, expected_queue_name in BOT_TYPE_TO_QUEUE_NAME.items():
            self.bot_profile.bot_type = bot_type
            self.bot_profile.save()

            self.second_bot_profile.bot_type = bot_type
            self.second_bot_profile.save()

            sender_email = self.user_profile.email
            recipient_emails = [self.bot_profile.email, self.second_bot_profile.email]
            profile_ids = [self.bot_profile.id, self.second_bot_profile.id]

            def check_values_passed(queue_name: Any,
                                    trigger_event: Union[Mapping[Any, Any], Any],
                                    x: Callable[[Any], None],
                                    call_consume_in_tests: bool) -> None:
                self.assertEqual(queue_name, expected_queue_name)
                self.assertIn(trigger_event["user_profile_id"], profile_ids)
                profile_ids.remove(trigger_event["user_profile_id"])
                self.assertEqual(trigger_event["trigger"], "private_message")
                self.assertEqual(trigger_event["message"]["sender_email"], sender_email)
                self.assertEqual(trigger_event["message"]["type"], u'private')
            mock_queue_json_publish.side_effect = check_values_passed

            self.send_huddle_message(sender_email, recipient_emails, 'test')
            self.assertEqual(mock_queue_json_publish.call_count, 2)
            mock_queue_json_publish.reset_mock()

    @mock.patch('zerver.lib.actions.queue_json_publish')
    def test_no_trigger_on_huddle_message_from_bot(self, mock_queue_json_publish: mock.Mock) -> None:
        for bot_type in BOT_TYPE_TO_QUEUE_NAME:
            self.bot_profile.bot_type = bot_type
            self.bot_profile.save()

            sender_email = self.second_bot_profile.email
            recipient_emails = [self.user_profile.email, self.bot_profile.email]
            self.send_huddle_message(sender_email, recipient_emails)
            self.assertFalse(mock_queue_json_publish.called)