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