JSON decode websocket events & allow usage of own loop

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.

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