refactoring

kvmd/apps/kvmd/server.py

@@ -57,7 +57,7 @@ from ...htserver import make_json_response
 from ...htserver import make_json_exception
 from ...htserver import send_ws_event
 from ...htserver import broadcast_ws_event
-from ...htserver import parse_ws_event
+from ...htserver import process_ws_messages
 from ...htserver import HttpServer
 
 from ...plugins import BasePlugin
@@ -278,21 +278,7 @@ class KvmdServer(HttpServer):  # pylint: disable=too-many-arguments,too-many-ins
                 ])
 
             await send_ws_event(ws, "loop", {})
-
+            await process_ws_messages(ws, self.__ws_handlers)
-            async for msg in ws:
-                if msg.type != aiohttp.web.WSMsgType.TEXT:
-                    break
-                try:
-                    (event_type, event) = parse_ws_event(msg.data)
-                except Exception as err:
-                    get_logger(0).error("Can't parse JSON event from websocket: %r", err)
-                else:
-                    handler = self.__ws_handlers.get(event_type)
-                    if handler:
-                        await handler(ws, event)
-                    else:
-                        get_logger(0).error("Unknown websocket event: %r", msg.data)
-
             return ws
         finally:
             await self.__remove_ws_client(client)

kvmd/htserver.py

@@ -38,6 +38,7 @@ from aiohttp.web import Request
 from aiohttp.web import Response
 from aiohttp.web import StreamResponse
 from aiohttp.web import WebSocketResponse
+from aiohttp.web import WSMsgType
 from aiohttp.web import Application
 from aiohttp.web import run_app
 from aiohttp.web import normalize_path_middleware
@@ -220,7 +221,7 @@ async def broadcast_ws_event(wss: List[WebSocketResponse], event_type: str, even
         ], return_exceptions=True)
 
 
-def parse_ws_event(msg: str) -> Tuple[str, Dict]:
+def _parse_ws_event(msg: str) -> Tuple[str, Dict]:
     data = json.loads(msg)
     if not isinstance(data, dict):
         raise RuntimeError("Top-level event structure is not a dict")
@@ -233,6 +234,23 @@ def parse_ws_event(msg: str) -> Tuple[str, Dict]:
     return (event_type, event)
 
 
+async def process_ws_messages(ws: WebSocketResponse, handlers: Dict[str, Callable]) -> None:
+    logger = get_logger(1)
+    async for msg in ws:
+        if msg.type != WSMsgType.TEXT:
+            break
+        try:
+            (event_type, event) = _parse_ws_event(msg.data)
+        except Exception as err:
+            logger.error("Can't parse JSON event from websocket: %r", err)
+        else:
+            handler = handlers.get(event_type)
+            if handler:
+                await handler(ws, event)
+            else:
+                logger.error("Unknown websocket event: %r", msg.data)
+
+
 # =====
 _REQUEST_AUTH_INFO = "_kvmd_auth_info"