diff --git a/src/mattermostdriver/driver.py b/src/mattermostdriver/driver.py index f45244c..ed4a0ac 100644 --- a/src/mattermostdriver/driver.py +++ b/src/mattermostdriver/driver.py @@ -333,12 +333,12 @@ def __enter__(self): def __exit__(self, *exc_info): return self.client.__exit__(*exc_info) - def init_websocket(self, event_handler, websocket_cls=Websocket): + def init_websocket(self, event_handler, websocket_cls=Websocket, loop=None): """ - Will initialize the websocket connection to the mattermost server. + Will initialize the websocket connection to the mattermost server and start an asyncio loop. This should be run after login(), because the websocket needs to make - an authentification. + an authentification. This function blocks until the connection with the server is broken. See https://api.mattermost.com/v4/#tag/WebSocket for which websocket events mattermost sends. @@ -353,13 +353,41 @@ async def my_event_handler(message): :param event_handler: The function to handle the websocket events. Takes one argument. :type event_handler: Function(message) + :param loop: The Asyncio loop to use. If left empty, get_event_loop is called. + :type loop: EventLoop, optional :return: The event loop """ self.websocket = websocket_cls(self.options, self.client.token) - loop = asyncio.get_event_loop() - loop.run_until_complete(self.websocket.connect(event_handler)) + if loop == None: + loop = asyncio.get_event_loop() + loop.run_until_complete(self.run_websocket(event_handler)) return loop + async def run_websocket(self, event_handler, websocket_cls=Websocket): + """ + Will initialize the websocket connection to the mattermost server and awaits messages. + + This should be run after login(), because the websocket needs to make + an authentification. + + See https://api.mattermost.com/v4/#tag/WebSocket for which + websocket events mattermost sends. + + Example of a really simple event_handler function + + .. code:: python + + async def my_event_handler(message): + print(message) + + + :param event_handler: The function to handle the websocket events. Takes one argument. + :type event_handler: Function(message) + :return: The event loop + """ + self.websocket = websocket_cls(self.options, self.client.token) + await self.websocket.connect(event_handler) + def login(self): """ Logs the user in. diff --git a/src/mattermostdriver/websocket.py b/src/mattermostdriver/websocket.py index 4576835..c38151e 100644 --- a/src/mattermostdriver/websocket.py +++ b/src/mattermostdriver/websocket.py @@ -3,7 +3,6 @@ import asyncio import logging import time - import aiohttp log = logging.getLogger("mattermostdriver.websocket") @@ -81,7 +80,9 @@ async def _start_loop(self, websocket, event_handler): while self._alive: message = await websocket.receive_str() self._last_msg = time.time() - await event_handler(message) + decoded = json.loads(message) + if "event" in decoded: + await event_handler(decoded) log.debug("cancelling heartbeat task") keep_alive.cancel() try: @@ -127,7 +128,8 @@ async def _authenticate_websocket(self, websocket, event_handler): log.debug(status) # We want to pass the events to the event_handler already # because the hello event could arrive before the authentication ok response - await event_handler(message) + if "event" in status: + await event_handler(status) if ("event" in status and status["event"] == "hello") and ("seq" in status and status["seq"] == 0): log.info("Websocket authentification OK") return True