big cleanup and refactoring #2: scramble master ws uri again

[hublib] / hub / websocket.py

diff --git a/hub/websocket.py b/hub/websocket.py
index 0ecf9875e29a5ba5f66665a1b54ef30bd2a41f9a..89a7e2f3cfa4fd79baed82bc75c6581620eccab5 100644 (file)
--- a/hub/websocket.py
+++ b/hub/websocket.py
@@ -9,24 +9,24 @@ from traceback import print_exception
 from falcon import WebSocketDisconnected
 
 
-class BaseWebSocketApp:
+class WebSocketApp:
     def __init__(self, hubapp):
-        self.hubapp = hubapp
-        self.conn = self.hubapp.app.hubapps["root"].conn
-
-    def task_done(self):
-        self.task = None
+        self.name = hubapp.name
+        self.conn = hubapp.root.conn
 
     @staticmethod
-    async def process_websocket(conn, web_socket, extra_data={}, recipients=[]):
+    async def process_websocket(conn, web_socket, extra_data=None, recipients=None):
         try:
             while True:
                 data = json.loads(await web_socket.receive_text())
-                data.update(extra_data)
+                if extra_data:
+                    data.update(extra_data)
                 if callable(recipients):
                     current_recipients = recipients(data)
-                else:
+                elif recipients:
                     current_recipients = recipients
+                else:
+                    raise ValueError("no recipients specified")
                 for recipient in current_recipients:
                     await conn.publish(recipient, pickle.dumps(data))
         except (CancelledError, WebSocketDisconnected):
@@ -76,11 +76,9 @@ class BaseWebSocketApp:
             if callable(leave_cb):
                 await leave_cb()
 
-
-class WebSocketApp(BaseWebSocketApp):
-    async def join_leave_client_notify(self, redis, action, client_id):
+    async def client_notify(self, redis, action, client_id):
         await redis.publish(
-            f"{self.hubapp.name}-master",
+            f"{self.name}-master",
             pickle.dumps({"action": action, "client_id": client_id}),
         )
 
@@ -89,25 +87,25 @@ class WebSocketApp(BaseWebSocketApp):
         return await self.on_websocket(
             req,
             web_socket,
-            f"{self.hubapp.name}-client-{client_id}",
+            f"{self.name}-client-{client_id}",
             {
                 "extra_data": {"client_id": client_id},
-                "recipients": [f"{self.hubapp.name}-master"],
+                "recipients": [f"{self.name}-master"],
             },
-            partial(self.join_leave_client_notify, self.conn, "join", client_id),
-            partial(self.join_leave_client_notify, self.conn, "leave", client_id),
+            partial(self.client_notify, self.conn, "join", client_id),
+            partial(self.client_notify, self.conn, "leave", client_id),
         )
 
     async def on_websocket_master(self, req, web_socket):
             return await self.on_websocket(
                 req,
                 web_socket,
-                f"{self.hubapp.name}-master",
+                f"{self.name}-master",
                 {"recipients": self.get_master_recipients},
             )
 
     def get_master_recipients(self, data):
         return [
-            f"{self.hubapp.name}-client-{int(client_id)}"
+            f"{self.name}-client-{int(client_id)}"
             for client_id in data.pop("client_ids", ())
         ]