Example #1
0
    def __init__(self,
                 settings=Settings(),
                 plugins=[ExamplePlugin(), WebHookExample()]):
        logging.basicConfig(
            **{
                "format": "[%(asctime)s] %(message)s",
                "datefmt": "%m/%d/%Y %H:%M:%S",
                "level": logging.DEBUG if settings.DEBUG else logging.INFO,
                "stream": sys.stdout,
            })
        self.settings = settings
        self.driver = Driver({
            "url": settings.MATTERMOST_URL,
            "port": settings.MATTERMOST_PORT,
            "token": settings.BOT_TOKEN,
            "scheme": settings.SCHEME,
            "verify": settings.SSL_VERIFY,
        })
        self.driver.login()
        self.plugins = self._initialize_plugins(plugins)
        self.event_handler = EventHandler(self.driver,
                                          settings=self.settings,
                                          plugins=self.plugins)
        self.webhook_server = None

        if self.settings.WEBHOOK_HOST_ENABLED:
            self._initialize_webhook_server()
Example #2
0
    def test_init(self):
        handler = EventHandler(Driver(),
                               Settings(),
                               plugins=[ExamplePlugin(),
                                        WebHookExample()])
        # Test the name matcher regexp
        assert handler._name_matcher.match("@my_username are you there?")
        assert not handler._name_matcher.match(
            "@other_username are you there?")

        # Test that all listeners from the individual plugins are now registered on
        # the handler
        for plugin in handler.plugins:
            for pattern, listener in plugin.message_listeners.items():
                assert listener in handler.message_listeners[pattern]
            for pattern, listener in plugin.webhook_listeners.items():
                assert listener in handler.webhook_listeners[pattern]

        # And vice versa, check that any listeners on the handler come from the
        # registered plugins
        for pattern, listeners in handler.message_listeners.items():
            for listener in listeners:
                assert any([
                    pattern in plugin.message_listeners
                    and listener in plugin.message_listeners[pattern]
                    for plugin in handler.plugins
                ])
        for pattern, listeners in handler.webhook_listeners.items():
            for listener in listeners:
                assert any([
                    pattern in plugin.webhook_listeners
                    and listener in plugin.webhook_listeners[pattern]
                    for plugin in handler.plugins
                ])
Example #3
0
    def test_handle_webhook(self):
        # Create an initialized plugin so its listeners are registered
        driver = Driver()
        plugin = WebHookExample().initialize(driver, Settings())
        # Construct a handler with it
        handler = EventHandler(driver, Settings(), plugins=[plugin])

        # Mock the call_function of the plugin so we can make some asserts
        async def mock_call_function(function, event, groups):
            # This is the regexp that we're trying to trigger
            assert function.matcher.pattern == "ping"
            assert event.text == "hello!"
            assert groups == []

        with mock.patch.object(plugin,
                               "call_function",
                               wraps=mock_call_function) as mocked:
            asyncio.run(
                handler._handle_webhook(
                    WebHookEvent(
                        body={"text": "hello!"},
                        request_id="request_id",
                        webhook_id="ping",
                    ), ))
            # Assert the function was called, so we know the asserts succeeded.
            mocked.assert_called_once()
Example #4
0
    def test_handle_post(self):
        # Create an initialized plugin so its listeners are registered
        driver = Driver()
        plugin = ExamplePlugin().initialize(driver)
        # Construct a handler with it
        handler = EventHandler(driver, Settings(), plugins=[plugin])

        # Mock the call_function of the plugin so we can make some asserts
        async def mock_call_function(function, message, groups):
            # This is the regexp that we're trying to trigger
            assert function.matcher.pattern == "sleep ([0-9]+)"
            assert message.text == "sleep 5"  # username should be stripped off
            assert groups == [
                "5"
            ]  # arguments should be matched and passed explicitly

        with mock.patch.object(plugin,
                               "call_function",
                               wraps=mock_call_function) as mocked:
            # Transform the default message into a raw post event so we can pass it
            new_body = create_message(text="@my_username sleep 5").body.copy()
            new_body["data"]["post"] = json.dumps(new_body["data"]["post"])
            new_body["data"]["mentions"] = json.dumps(
                new_body["data"]["mentions"])
            asyncio.run(handler._handle_post(new_body))

            # Assert the function was called, so we know the asserts succeeded.
            mocked.assert_called_once()
Example #5
0
    def test_handle_event(self, handle_post):
        handler = EventHandler(Driver(), Settings(), plugins=[])
        # This event should trigger _handle_post
        asyncio.run(handler._handle_event(json.dumps(create_message().body)))
        # This event should not
        asyncio.run(
            handler._handle_event(json.dumps({"event": "some_other_event"})))

        handle_post.assert_called_once_with(create_message().body)
Example #6
0
    def test_should_ignore(self):
        handler = EventHandler(Driver(),
                               Settings(IGNORE_USERS=["ignore_me"]),
                               plugins=[])
        # We shouldn't ignore a message from betty, since she is not listed
        assert not handler._should_ignore(create_message(sender_name="betty"))
        assert handler._should_ignore(create_message(sender_name="ignore_me"))

        # We ignore our own messages by default
        assert handler._should_ignore(
            create_message(sender_name="my_username"))

        # But shouldn't do so if this is explicitly requested
        handler = EventHandler(
            Driver(),
            Settings(IGNORE_USERS=["ignore_me"]),
            plugins=[],
            ignore_own_messages=False,
        )
        assert not handler._should_ignore(
            create_message(sender_name="my_username"))
Example #7
0
    def test_allowed_users(self):
        wrapped = mock.create_autospec(example_listener)
        wrapped.__qualname__ = "wrapped"
        # Create a driver with a mocked reply function
        driver = Driver()

        def fake_reply(message, text):
            assert "you do not have permission" in text.lower()

        driver.reply_to = mock.Mock(wraps=fake_reply)

        f = listen_to("", allowed_users=["Betty"])(wrapped)
        f.plugin = ExamplePlugin().initialize(driver)

        # This is fine, the names are not caps sensitive
        f(create_message(sender_name="betty"))
        wrapped.assert_called_once()
        wrapped.reset_mock()

        # This is not fine, and we expect the fake reply to be called.
        f(create_message(sender_name="not_betty"))
        wrapped.assert_not_called()
        driver.reply_to.assert_called_once()
Example #8
0
def expect_reply(driver: Driver, post: Dict, wait=RESPONSE_TIMEOUT, retries=1):
    """Utility function to specify we expect some kind of reply after `wait` seconds."""
    reply = None
    for _ in range(retries + 1):
        time.sleep(wait)
        thread_info = driver.get_thread(post["id"])
        print(thread_info)
        reply_id = thread_info["order"][-1]
        if reply_id != post["id"]:
            reply = thread_info["posts"][reply_id]
            break

    if not reply:
        raise ValueError("Expected a response, but didn't get any!")

    return reply
Example #9
0
    def test_needs_mention(self):  # noqa
        wrapped = mock.create_autospec(example_listener)
        wrapped.__qualname__ = "wrapped"
        f = listen_to("", needs_mention=True)(wrapped)
        f.plugin = ExamplePlugin().initialize(Driver())

        # The default message mentions the specified user ID, so should be called
        f(create_message(mentions=["qmw86q7qsjriura9jos75i4why"]))
        wrapped.assert_called_once()
        wrapped.reset_mock()

        # No mention, so the function should still only have been called once in total
        f(create_message(mentions=[]))
        wrapped.assert_not_called()

        # But if this is a direct message, we do want to trigger
        f(create_message(mentions=[], channel_type="D"))
        wrapped.assert_called_once()
Example #10
0
    def test_call_function(self, add_task):
        p = FakePlugin().initialize(Driver())

        # Since this is not an async function, a task should be added to the threadpool
        message = create_message(text="pattern")
        asyncio.run(
            p.call_function(FakePlugin.my_function,
                            message,
                            groups=["test", "another"]))
        add_task.assert_called_once_with(FakePlugin.my_function, message,
                                         "test", "another")

        # Since this is an async function, it should be called directly through asyncio.
        message = create_message(text="async_pattern")
        with mock.patch.object(p.my_async_function,
                               "function") as mock_function:
            asyncio.run(
                p.call_function(FakePlugin.my_async_function,
                                message,
                                groups=[]))
            mock_function.assert_called_once_with(p, message)
Example #11
0
    def test_initialize(self):
        p = FakePlugin().initialize(Driver())
        # Test whether the function was registered properly
        assert p.message_listeners[re.compile("pattern")] == [
            FakePlugin.my_function,
        ]

        # This function should be registered twice, once for each listener
        assert len(p.message_listeners[re.compile("async_pattern")]) == 1
        assert (p.message_listeners[re.compile("async_pattern")][0].function ==
                FakePlugin.my_async_function.function)

        assert len(
            p.message_listeners[re.compile("another_async_pattern")]) == 1
        assert (p.message_listeners[re.compile("another_async_pattern")]
                [0].function == FakePlugin.my_async_function.function)

        assert len(p.webhook_listeners) == 1
        assert p.webhook_listeners[re.compile("webhook_id")] == [
            FakePlugin.webhook_listener
        ]
Example #12
0
    def test_click_function(self):
        @click.command()
        @click.option("--arg1", type=str, default="nothing")
        @click.option("--arg2", type=str, default="nothing either")
        @click.option("-f", "--flag", is_flag=True)
        def wrapped(self, message, arg1, arg2, flag):
            return arg1, arg2, flag

        f = MessageFunction(wrapped, matcher=re.compile(""))
        # Verify that the arguments are passed and returned correctly
        assert f(create_message(),
                 "--arg1=yes --arg2=no") == ("yes", "no", False)
        assert f(create_message(), "-f --arg2=no") == ("nothing", "no", True)

        # If an incorrect argument is passed, the error and help string should be returned.
        def mocked_reply(message, response):
            assert "no such option: --nonexistent-arg" in response
            assert f.docstring in response

        f.plugin = ExamplePlugin().initialize(Driver(), Settings())
        with mock.patch.object(f.plugin.driver, "reply_to",
                               wraps=mocked_reply) as mock_function:
            f(create_message(), "-f --arg2=no --nonexistent-arg")
            mock_function.assert_called_once()
Example #13
0
    def test_ensure_response(self):
        p = ExamplePlugin().initialize(Driver())

        def mock_respond(event, response):
            event.responded = True

        # SCENARIO 1: We wrap a function that does not send a web response
        with mock.patch.object(p.driver, "respond_to_web",
                               wraps=mock_respond) as mocked:
            f = WebHookFunction(example_webhook_listener,
                                matcher=re.compile(""))
            f.plugin = p
            event = create_webhook_event()

            f(event)
            # We expect the WebHookFunction to automatically respond NoResponse
            mocked.assert_called_once_with(event, NoResponse)
            assert event.responded

        # SCECNARIO 2: Same scenario as above, but with asyncio
        async def webhook_function(self, event):
            pass

        with mock.patch.object(p.driver, "respond_to_web",
                               wraps=mock_respond) as mocked:
            f = WebHookFunction(webhook_function, matcher=re.compile(""))
            f.plugin = p
            event = create_webhook_event()

            # Asyncio helper to emulate running from async function
            async def run():
                await f(event)

            asyncio.run(run())
            # We expect the WebHookFunction to automatically respond NoResponse
            mocked.assert_called_once_with(event, NoResponse)
            assert event.responded

        # SCENARIO 3: We wrap a function that does send a web response
        def webhook_function(self, event):
            self.driver.respond_to_web(event, "Hello!")

        with mock.patch.object(p.driver, "respond_to_web",
                               wraps=mock_respond) as mocked:
            f = WebHookFunction(webhook_function, matcher=re.compile(""))
            f.plugin = p
            event = create_webhook_event()

            f(event)
            # We expect the WebHookFunction to not respond anything, since the function
            # itself already responded 'Hello!'.
            mocked.assert_called_once_with(event, "Hello!")
            assert event.responded

        # SCNEARIO 4: Same scenario as above, but with asyncio
        async def webhook_function(self, event):
            self.driver.respond_to_web(event, "Hello!")

        with mock.patch.object(p.driver, "respond_to_web",
                               wraps=mock_respond) as mocked:
            f = WebHookFunction(webhook_function, matcher=re.compile(""))
            f.plugin = p
            event = create_webhook_event()

            # Asyncio helper to emulate running from async function
            async def run():
                await f(event)

            asyncio.run(run())
            # We expect the WebHookFunction to not respond anything, since the function
            # itself already responded 'Hello!'.
            mocked.assert_called_once_with(event, "Hello!")
            assert event.responded
Example #14
0
class Bot:
    """Base chatbot class.

    Can be either subclassed for custom functionality, or used as-is with custom plugins
    and settings. To start the bot, simply call bot.run().
    """
    def __init__(self,
                 settings=Settings(),
                 plugins=[ExamplePlugin(), WebHookExample()]):
        logging.basicConfig(
            **{
                "format": "[%(asctime)s] %(message)s",
                "datefmt": "%m/%d/%Y %H:%M:%S",
                "level": logging.DEBUG if settings.DEBUG else logging.INFO,
                "stream": sys.stdout,
            })
        self.settings = settings
        self.driver = Driver({
            "url": settings.MATTERMOST_URL,
            "port": settings.MATTERMOST_PORT,
            "token": settings.BOT_TOKEN,
            "scheme": settings.SCHEME,
            "verify": settings.SSL_VERIFY,
        })
        self.driver.login()
        self.plugins = self._initialize_plugins(plugins)
        self.event_handler = EventHandler(self.driver,
                                          settings=self.settings,
                                          plugins=self.plugins)
        self.webhook_server = None

        if self.settings.WEBHOOK_HOST_ENABLED:
            self._initialize_webhook_server()

    def _initialize_plugins(self, plugins: Sequence[Plugin]):
        for plugin in plugins:
            plugin.initialize(self.driver, self.settings)
        return plugins

    def _initialize_webhook_server(self):
        self.webhook_server = WebHookServer(
            url=self.settings.WEBHOOK_HOST_URL,
            port=self.settings.WEBHOOK_HOST_PORT)
        self.driver.register_webhook_server(self.webhook_server)
        # Schedule the queue loop to the current event loop so that it starts together
        # with self.init_websocket.
        asyncio.get_event_loop().create_task(
            self.event_handler._check_queue_loop(
                self.webhook_server.event_queue))

    def run(self):
        logging.info(f"Starting bot {self.__class__.__name__}.")
        try:
            self.driver.threadpool.start()
            # Start a thread to run potential scheduled jobs
            self.driver.threadpool.start_scheduler_thread(
                self.settings.SCHEDULER_PERIOD)
            # Start the webhook server on a separate thread if necessary
            if self.settings.WEBHOOK_HOST_ENABLED:
                self.driver.threadpool.start_webhook_server_thread(
                    self.webhook_server)

            for plugin in self.plugins:
                plugin.on_start()

            # Start listening for events
            self.event_handler.start()

        except KeyboardInterrupt as e:
            self.stop()
            raise e

    def stop(self):
        logging.info("Stopping bot.")
        # Shutdown the running plugins
        for plugin in self.plugins:
            plugin.on_stop()
        # Stop the threadpool
        self.driver.threadpool.stop()
Example #15
0
 def test_help_string(self, snapshot):
     p = FakePlugin().initialize(Driver())
     # Compare the help string with the snapshotted version.
     snapshot.assert_match(p.get_help_string())