From 35785be5dc08e82f827336bbc1a2629b5579268e Mon Sep 17 00:00:00 2001
From: Eric Callahan <arksine.code@gmail.com>
Date: Sat, 13 Jan 2024 19:47:08 -0500
Subject: [PATCH] application: remove direct websockets import

Provide a method for websocket handler registration.

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
---
 moonraker/components/application.py | 18 +++++++++++-------
 moonraker/components/websockets.py  |  9 ++++++++-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/moonraker/components/application.py b/moonraker/components/application.py
index 9c45d68..ab32b95 100644
--- a/moonraker/components/application.py
+++ b/moonraker/components/application.py
@@ -33,7 +33,6 @@ from ..common import (
     KlippyState
 )
 from ..utils import json_wrapper as jsonw
-from .websockets import WebSocket, BridgeSocket
 from streaming_form_data import StreamingFormDataParser, ParseFailedException
 from streaming_form_data.targets import FileTarget, ValueTarget, SHA256Target
 
@@ -48,15 +47,17 @@ from typing import (
     Dict,
     List,
     AsyncGenerator,
+    Type
 )
 if TYPE_CHECKING:
     from tornado.httpserver import HTTPServer
+    from tornado.websocket import WebSocketHandler
     from ..server import Server
     from ..eventloop import EventLoop
     from ..confighelper import ConfigHelper
     from .klippy_connection import KlippyConnection as Klippy
     from ..utils import IPAddress
-    from .websockets import WebsocketManager
+    from .websockets import WebsocketManager, WebSocket
     from .file_manager.file_manager import FileManager
     from .announcements import Announcements
     from .machine import Machine
@@ -196,8 +197,6 @@ class MoonrakerApp:
         app_handlers: List[Any] = [
             (AnyMatches(), self.mutable_router),
             (home_pattern, WelcomeHandler),
-            (f"{self._route_prefix}/websocket", WebSocket),
-            (f"{self._route_prefix}/klippysocket", BridgeSocket),
             (f"{self._route_prefix}/server/redirect", RedirectHandler),
             (f"{self._route_prefix}/server/jsonrpc", RPCHandler)
         ]
@@ -375,6 +374,13 @@ class MoonrakerApp:
             f"{self._route_prefix}{pattern}", FileUploadHandler, params
         )
 
+    def register_websocket_handler(
+        self, pattern: str, handler: Type[WebSocketHandler]
+    ) -> None:
+        self.mutable_router.add_handler(
+            f"{self._route_prefix}{pattern}", handler, None
+        )
+
     def register_debug_endpoint(
         self,
         endpoint: str,
@@ -459,9 +465,7 @@ class AuthorizedRequestHandler(tornado.web.RequestHandler):
                 pass
             else:
                 wsm: WebsocketManager = self.server.lookup_component("websockets")
-                conn = wsm.get_client(conn_id)
-        if not isinstance(conn, WebSocket):
-            return None
+                conn = wsm.get_client_ws(conn_id)
         return conn
 
     def write_error(self, status_code: int, **kwargs) -> None:
diff --git a/moonraker/components/websockets.py b/moonraker/components/websockets.py
index f83a467..bd14460 100644
--- a/moonraker/components/websockets.py
+++ b/moonraker/components/websockets.py
@@ -34,6 +34,7 @@ if TYPE_CHECKING:
     from ..server import Server
     from .klippy_connection import KlippyConnection as Klippy
     from ..confighelper import ConfigHelper
+    from .application import MoonrakerApp
     from .extensions import ExtensionManager
     from .authorization import Authorization
     from ..utils import IPAddress
@@ -50,6 +51,9 @@ class WebsocketManager:
         self.clients: Dict[int, BaseRemoteConnection] = {}
         self.bridge_connections: Dict[int, BridgeSocket] = {}
         self.closed_event: Optional[asyncio.Event] = None
+        app: MoonrakerApp = self.server.lookup_component("application")
+        app.register_websocket_handler("/websocket", WebSocket)
+        app.register_websocket_handler("/klippysocket", BridgeSocket)
         self.server.register_endpoint(
             "/server/websocket/id", RequestType.GET, self._handle_id_request,
             TransportType.WEBSOCKET
@@ -130,7 +134,10 @@ class WebsocketManager:
     def has_socket(self, ws_id: int) -> bool:
         return ws_id in self.clients
 
-    def get_client(self, ws_id: int) -> Optional[BaseRemoteConnection]:
+    def get_client(self, uid: int) -> Optional[BaseRemoteConnection]:
+        return self.clients.get(uid, None)
+
+    def get_client_ws(self, ws_id: int) -> Optional[WebSocket]:
         sc = self.clients.get(ws_id, None)
         if sc is None or not isinstance(sc, WebSocket):
             return None