refactoring

2 files changed, 11 insertions, 21 deletions

diff --git a/kvmd/apps/kvmd/__init__.py b/kvmd/apps/kvmd/__init__.py
index 3de764fd..f7536577 100644
--- a/kvmd/apps/kvmd/__init__.py
+++ b/kvmd/apps/kvmd/__init__.py
@@ -100,14 +100,12 @@ def main(argv: Optional[List[str]]=None) -> None:
             **config.snapshot._unpack(),
         ),
 
-        heartbeat=config.server.heartbeat,
-
         keymap_path=config.hid.keymap,
         ignore_keys=config.hid.ignore_keys,
         mouse_x_range=(config.hid.mouse_x_range.min, config.hid.mouse_x_range.max),
         mouse_y_range=(config.hid.mouse_y_range.min, config.hid.mouse_y_range.max),
 
         stream_forever=config.streamer.forever,
-    ).run(**config.server._unpack(ignore=["heartbeat"]))
+    ).run(**config.server._unpack())
 
     get_logger(0).info("Bye-bye")
diff --git a/kvmd/apps/kvmd/server.py b/kvmd/apps/kvmd/server.py
index e7bc744d..f40381a9 100644
--- a/kvmd/apps/kvmd/server.py
+++ b/kvmd/apps/kvmd/server.py
@@ -157,8 +157,6 @@ class KvmdServer(HttpServer):  # pylint: disable=too-many-arguments,too-many-ins
         streamer: Streamer,
         snapshoter: Snapshoter,
 
-        heartbeat: float,
-
         keymap_path: str,
         ignore_keys: List[str],
         mouse_x_range: Tuple[int, int],
@@ -173,8 +171,6 @@ class KvmdServer(HttpServer):  # pylint: disable=too-many-arguments,too-many-ins
         self.__snapshoter = snapshoter  # Not a component: No state or cleanup
         self.__user_gpio = user_gpio  # Has extra state "gpio_scheme_state"
 
-        self.__heartbeat = heartbeat
-
         self.__stream_forever = stream_forever
 
         self.__components = [
@@ -254,13 +250,9 @@ class KvmdServer(HttpServer):  # pylint: disable=too-many-arguments,too-many-ins
 
     @exposed_http("GET", "/ws")
     async def __ws_handler(self, request: aiohttp.web.Request) -> aiohttp.web.WebSocketResponse:
-        logger = get_logger(0)
-
-        client = _WsClient(
-            ws=aiohttp.web.WebSocketResponse(heartbeat=self.__heartbeat),
-            stream=valid_bool(request.query.get("stream", "true")),
-        )
-        await client.ws.prepare(request)
+        stream = valid_bool(request.query.get("stream", "true"))
+        ws = await self._make_ws_response(request)
+        client = _WsClient(ws, stream)
         await self.__register_ws_client(client)
 
         try:
@@ -281,27 +273,27 @@ class KvmdServer(HttpServer):  # pylint: disable=too-many-arguments,too-many-ins
             ))
             for stage in [stage1, stage2]:
                 await asyncio.gather(*[
-                    send_ws_event(client.ws, event_type, events.pop(event_type))
+                    send_ws_event(ws, event_type, events.pop(event_type))
                     for (event_type, _) in stage
                 ])
 
-            await send_ws_event(client.ws, "loop", {})
+            await send_ws_event(ws, "loop", {})
 
-            async for msg in client.ws:
+            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:
-                    logger.error("Can't parse JSON event from websocket: %r", 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(client.ws, event)
+                        await handler(ws, event)
                     else:
-                        logger.error("Unknown websocket event: %r", msg.data)
+                        get_logger(0).error("Unknown websocket event: %r", msg.data)
 
-            return client.ws
+            return ws
         finally:
             await self.__remove_ws_client(client)