From 12487cdb645e4cf667e915d2064d308c01f1cde8 Mon Sep 17 00:00:00 2001 From: Claudiu Date: Thu, 23 Jun 2016 21:39:29 -0700 Subject: [PATCH] Add bot-initiated messaging ability The approach is to use an idle handler, which gets called whenever something else hasn't already happened, or once per second. This should help reduce rate limiting. Functional testing has been added. Updated README.md with how to use the new functionality. Resolves #92. --- .gitignore | 2 + README.md | 23 +++++++-- slackbot/bot.py | 12 +++++ slackbot/dispatcher.py | 32 +++++++++++- slackbot/manager.py | 6 ++- slackbot/plugins/hello.py | 51 +++++++++++++++++-- slackbot/slackclient.py | 79 +++++++++++++++++++++++++++-- slackbot/utils.py | 17 +++++++ tests/functional/test_functional.py | 44 ++++++++++++++++ 9 files changed, 253 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index d3e1bedc..ca5fb4e9 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,5 @@ slackbot_test_settings.py /dist /*.egg-info .cache +.idea + diff --git a/README.md b/README.md index 0052ae0a..aac24c6a 100644 --- a/README.md +++ b/README.md @@ -108,14 +108,14 @@ def github(): A chat bot is meaningless unless you can extend/customize it to fit your own use cases. -To write a new plugin, simplely create a function decorated by `slackbot.bot.respond_to` or `slackbot.bot.listen_to`: +To write a new plugin, simply create a function decorated by `slackbot.bot.respond_to`, `slackbot.bot.listen_to`, or `slackbot.bot.idle`: - A function decorated with `respond_to` is called when a message matching the pattern is sent to the bot (direct message or @botname in a channel/group chat) - A function decorated with `listen_to` is called when a message matching the pattern is sent on a channel/group chat (not directly sent to the bot) +- *(development version only)* A function decorated with `idle` is called whenever a message has not been sent for the past second ```python -from slackbot.bot import respond_to -from slackbot.bot import listen_to +from slackbot.bot import respond_to, listen_to, idle import re @respond_to('hi', re.IGNORECASE) @@ -135,6 +135,23 @@ def help(message): # Message is sent on the channel # message.send('I can help everybody!') + +last_bored = time.time() +@idle +def bored(client): + if time.time() - last_bored >= 30: + last_bored = time.time() + + # Messages can be sent to a channel + client.rtm_send_message('some_channel', "I'm bored!") + # Or directly to a user + client.rtm_send_message('some_user', "Hey, entertain me!") + + # If a name is ambiguous: + client.rtm_send_message(client.find_channel_by_name('ambiguous'), "To ambiguous the channel") + client.rtm_send_message(client.find_user_by_name('ambiguous'), "To ambiguous the user") + + # Attachments can be sent with `client.rtm_send_message(..., attachments=attachments)`. ``` To extract params from the message, you can use regular expression: diff --git a/slackbot/bot.py b/slackbot/bot.py index a45010af..39592678 100644 --- a/slackbot/bot.py +++ b/slackbot/bot.py @@ -11,6 +11,7 @@ from slackbot.manager import PluginsManager from slackbot.slackclient import SlackClient from slackbot.dispatcher import MessageDispatcher +from slackbot.utils import optional_arg_decorator logger = logging.getLogger(__name__) @@ -65,6 +66,17 @@ def wrapper(func): return wrapper +# use optional_arg_decorator so users can either do @idle or @idle() +@optional_arg_decorator +def idle(func): + """Run a function once/second whenever no other actions were taken. + The function must take one parameter, a SlackClient instance.""" + # match anything, the text doesn't apply for "idle" + PluginsManager.idle_commands.append(func) + logger.info('registered idle plugin "%s"', func.__name__) + return func + + # def default_reply(matchstr=r'^.*$', flags=0): def default_reply(*args, **kwargs): """ diff --git a/slackbot/dispatcher.py b/slackbot/dispatcher.py index 14f847e7..ff68e1ac 100644 --- a/slackbot/dispatcher.py +++ b/slackbot/dispatcher.py @@ -135,13 +135,43 @@ def filter_text(self, msg): return msg def loop(self): + # once/second, check events + # run idle handlers whenever idle while True: events = self._client.rtm_read() for event in events: if event.get('type') != 'message': continue self._on_new_message(event) - time.sleep(1) + + # run idle handlers as long as we've been idle + for func in self._plugins.get_idle_plugins(): + if not func: + continue + + # if actions are pending, don't do anything + if not self._pool.queue.empty(): + break + + # if some action was taken, don't run the remaining handlers + if self._client.idle_time() < 1: + break + + try: + func(self._client) + except: + logger.exception( + 'idle handler failed with plugin "%s"', + func.__name__) + reply = u'[{}] I had a problem with idle handler\n'.format( + func.__name__) + tb = u'```\n{}\n```'.format(traceback.format_exc()) + # no channel, so only send errors to error user + if self._errors_to: + self._client.rtm_send_message(self._errors_to, + '{}\n{}'.format(reply, + tb)) + time.sleep(1.0) def _default_reply(self, msg): default_reply = settings.DEFAULT_REPLY diff --git a/slackbot/manager.py b/slackbot/manager.py index 0239e238..48859358 100644 --- a/slackbot/manager.py +++ b/slackbot/manager.py @@ -18,8 +18,9 @@ def __init__(self): commands = { 'respond_to': {}, 'listen_to': {}, - 'default_reply': {} + 'default_reply': {}, } + idle_commands = [] def init_plugins(self): if hasattr(settings, 'PLUGINS'): @@ -72,3 +73,6 @@ def get_plugins(self, category, text): if not has_matching_plugin: yield None, None + + def get_idle_plugins(self): + yield from self.idle_commands diff --git a/slackbot/plugins/hello.py b/slackbot/plugins/hello.py index d2b81937..1b7b0893 100644 --- a/slackbot/plugins/hello.py +++ b/slackbot/plugins/hello.py @@ -1,7 +1,8 @@ -#coding: UTF-8 +# coding: UTF-8 +import random import re -from slackbot.bot import respond_to -from slackbot.bot import listen_to + +from slackbot.bot import respond_to, listen_to, idle @respond_to('hello$', re.IGNORECASE) @@ -52,3 +53,47 @@ def hey(message): @respond_to(u'你好') def hello_unicode_message(message): message.reply(u'你好!') + + +# idle tests +IDLE_TEST = {'which': None, 'channel': None} + + +@respond_to('start idle test ([0-9]+)') +@listen_to('start idle test ([0-9]+)') +def start_idle_test(message, i): + print("---------- start idle test! -----------") + IDLE_TEST['which'] = int(i) + IDLE_TEST['channel'] = message._body['channel'] + print("Idle test is now {which} on channel {channel}".format(**IDLE_TEST)) + # TESTING ONLY, don't rely on this behavior + + +# idle function testing +# tests 0 and 1: rtm and webapi work from idle function 1 +# tests 2 and 3: rtm and webapi work from idle function 2 +# test 4: both idle functions can operate simultaneously +@idle +def idle_1(client): + which = IDLE_TEST['which'] + msg = "I am bored %s" % which + if which == 0: + client.rtm_send_message(IDLE_TEST['channel'], msg) + elif which == 1: + client.send_message(IDLE_TEST['channel'], msg) + elif which == 4: + if random.random() <= 0.5: + client.rtm_send_message(IDLE_TEST['channel'], "idle_1 is bored") + + +@idle() +def idle_2(client): + which = IDLE_TEST['which'] + msg = "I am bored %s" % which + if which == 2: + client.rtm_send_message(IDLE_TEST['channel'], msg) + elif which == 3: + client.send_message(IDLE_TEST['channel'], msg) + elif which == 4: + if random.random() <= 0.5: + client.rtm_send_message(IDLE_TEST['channel'], "idle_2 is bored") diff --git a/slackbot/slackclient.py b/slackbot/slackclient.py index 6c10992f..f180d65e 100644 --- a/slackbot/slackclient.py +++ b/slackbot/slackclient.py @@ -30,9 +30,13 @@ def __init__(self, token, bot_icon=None, bot_emoji=None, connect=True): self.websocket = None self.users = {} self.channels = {} + self.dm_channels = {} # map user id to direct message channel id self.connected = False self.webapi = slacker.Slacker(self.token) + # keep track of last action for idle handling + self._last_action = time.time() + if connect: self.rtm_connect() @@ -65,14 +69,22 @@ def parse_slack_login_data(self, login_data): def parse_channel_data(self, channel_data): self.channels.update({c['id']: c for c in channel_data}) + # pre-load direct message channels + for c in channel_data: + if 'user' in c: + self.dm_channels[c['user']] = c['id'] def send_to_websocket(self, data): - """Send (data) directly to the websocket.""" + """Send (data) directly to the websocket. + + Update last action for idle handling.""" data = json.dumps(data) self.websocket.send(data) + self._last_action = time.time() def ping(self): - return self.send_to_websocket({'type': 'ping'}) + self.send_to_websocket({'type': 'ping'}) + self._last_action = time.time() def websocket_safe_read(self): """Returns data if available, otherwise ''. Newlines indicate multiple messages """ @@ -101,7 +113,8 @@ def rtm_read(self): data.append(json.loads(d)) return data - def rtm_send_message(self, channel, message, attachments=None): + def rtm_send_message(self, channelish, message, attachments=None): + channel = self._channelify(channelish) message_json = { 'type': 'message', 'channel': channel, @@ -110,14 +123,17 @@ def rtm_send_message(self, channel, message, attachments=None): } self.send_to_websocket(message_json) - def upload_file(self, channel, fname, fpath, comment): + def upload_file(self, channelish, fname, fpath, comment): + channel = self._channelify(channelish) fname = fname or to_utf8(os.path.basename(fpath)) self.webapi.files.upload(fpath, channels=channel, filename=fname, initial_comment=comment) + self._last_action = time.time() - def send_message(self, channel, message, attachments=None, as_user=True): + def send_message(self, channelish, message, attachments=None, as_user=True): + channel = self._channelify(channelish) self.webapi.chat.post_message( channel, message, @@ -126,10 +142,57 @@ def send_message(self, channel, message, attachments=None, as_user=True): icon_emoji=self.bot_emoji, attachments=attachments, as_user=as_user) + self._last_action = time.time() def get_channel(self, channel_id): return Channel(self, self.channels[channel_id]) + def get_dm_channel(self, user_id): + """Get the direct message channel for the given user id, opening + one if necessary.""" + if user_id not in self.users: + raise ValueError("Expected valid user_id, have no user '%s'" % ( + user_id,)) + + if user_id in self.dm_channels: + return self.dm_channels[user_id] + + # open a new channel + resp = self.webapi.im.open(user_id) + if not resp.body["ok"]: + raise ValueError("Could not open DM channel: %s" % resp.body) + + self.dm_channels[user_id] = resp.body['channel']['id'] + + return self.dm_channels[user_id] + + def _channelify(self, s): + """Turn a string into a channel. + + * Given a channel id, return that same channel id. + * Given a channel name, return the channel id. + * Given a user id, return the direct message channel with that user, + opening a new one if necessary. + * Given a user name, do the same as for a user id. + + Raise a ValueError otherwise.""" + if s in self.channels: + return s + + channel_id = self.find_channel_by_name(s) + if channel_id: + return channel_id + + if s in self.users: + return self.get_dm_channel(s) + + user_id = self.find_user_by_name(s) + if user_id: + return self.get_dm_channel(user_id) + + raise ValueError("Could not turn '%s' into any kind of channel name" % ( + user_id)) + def find_channel_by_name(self, channel_name): for channel_id, channel in iteritems(self.channels): try: @@ -149,6 +212,12 @@ def react_to_message(self, emojiname, channel, timestamp): name=emojiname, channel=channel, timestamp=timestamp) + self._last_action = time.time() + + def idle_time(self): + """Return the time the client has been idle, i.e. the time since + it sent the last message to the server.""" + return time.time() - self._last_action class SlackConnectionError(Exception): diff --git a/slackbot/utils.py b/slackbot/utils.py index b79911c8..26cf206e 100644 --- a/slackbot/utils.py +++ b/slackbot/utils.py @@ -76,3 +76,20 @@ def do_work(self): while True: msg = self.queue.get() self.func(msg) + + +def optional_arg_decorator(fn): + """Allows for easier making of decorators with optional arguments. + + See: http://stackoverflow.com/questions/3888158/python-making-decorators-with-optional-arguments""" + def wrapped_decorator(*args): + if len(args) == 1 and callable(args[0]): + return fn(args[0]) + + else: + def real_decorator(decoratee): + return fn(decoratee, *args) + + return real_decorator + + return wrapped_decorator diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py index 1906efda..d2a24a3f 100644 --- a/tests/functional/test_functional.py +++ b/tests/functional/test_functional.py @@ -8,7 +8,9 @@ import os from os.path import join, abspath, dirname, basename import subprocess + import pytest + from tests.functional.driver import Driver from tests.functional.slackbot_settings import ( testbot_apitoken, testbot_username, @@ -192,3 +194,45 @@ def test_bot_reply_with_alias_message(driver): driver.wait_for_bot_channel_message("hello sender!", tosender=True) driver.send_channel_message('!hello', tobot=False, colon=False) driver.wait_for_bot_channel_message("hello sender!", tosender=True) + + +def make_idle_func(use_rtm): + def idle_func_bored(client): + if use_rtm: + client.rtm_send_message(test_channel, "I am bored") + else: + client.send_message(test_channel, "I am bored") + return idle_func_bored + + +# 5 tests, defined in hello.py +# parametrize based on method for fewer test functions +send_methods = { + 'direct': { + 'send': lambda driver, msg: driver.send_direct_message(msg), + 'wait': lambda driver, msg: driver.wait_for_bot_direct_message(msg), + }, + 'channel': { + 'send': lambda driver, msg: driver.send_channel_message(msg), + 'wait': lambda driver, msg: driver.wait_for_bot_channel_message(msg, tosender=False) + }, + 'group': { + 'send': lambda driver, msg: driver.send_group_message(msg), + 'wait': lambda driver, msg: driver.wait_for_bot_group_message(msg, tosender=False) + }, +} + +@pytest.mark.parametrize('which_test', [0, 1, 2, 3]) +@pytest.mark.parametrize('method', send_methods.keys()) +def test_idle_func_dm(driver, which_test, method): + send_methods[method]['send'](driver, "start idle test %d" % which_test) + send_methods[method]['wait'](driver, "I am bored %d" % which_test) + + +@pytest.mark.parametrize('method', send_methods.keys()) +def test_idle_func_dm_multi(driver, method): + send_methods[method]['send'](driver, "start idle test 4") + send_methods[method]['wait'](driver, "idle_1 is bored") + send_methods[method]['wait'](driver, "idle_2 is bored") + send_methods[method]['wait'](driver, "idle_1 is bored") + send_methods[method]['wait'](driver, "idle_2 is bored")