1 files changed, 31 insertions, 17 deletions

diff --git a/proxy/websockets.py b/proxy/websockets.py
index ec65e6d..a09e261 100644
--- a/proxy/websockets.py
+++ b/proxy/websockets.py
@@ -9,17 +9,22 @@ from uuid import uuid4
 from starlette.endpoints import WebSocketEndpoint
 from starlette.websockets import WebSocket
 
-from proxy.utils import get_path_with_query_string
 from proxy.queue import ProxyQueue
-from ttun_server.types import WebsocketMessage, WebsocketMessageType, WebsocketConnectData, WebsocketMessageData, \
-    WebsocketDisconnectData
+from proxy.utils import get_path_with_query_string
+from ttun_server.types import (
+    WebsocketConnectData,
+    WebsocketDisconnectData,
+    WebsocketMessage,
+    WebsocketMessageData,
+    WebsocketMessageType,
+)
 
 logger = logging.getLogger(__name__)
-logger.setLevel('DEBUG')
+logger.setLevel("DEBUG")
 
 
 class WebsocketProxy(WebSocketEndpoint):
-    encoding = 'json'
+    encoding = "json"
     websocket_listen_task = None
 
     def __init__(self, *args, **kwargs):
@@ -28,16 +33,20 @@ class WebsocketProxy(WebSocketEndpoint):
 
     @asynccontextmanager
     async def proxy(self, websocket: WebSocket, message: WebsocketMessage):
-        [subdomain, *_] = websocket.url.hostname.split('.')
+        [subdomain, *_] = websocket.url.hostname.split(".")
 
-        expect_ack = WebsocketMessageType(message['type']) == WebsocketMessageType.connect
+        expect_ack = (
+            WebsocketMessageType(message["type"]) == WebsocketMessageType.connect
+        )
 
         try:
             request_queue = await ProxyQueue.get_for_identifier(subdomain)
             await request_queue.enqueue(message)
 
             if expect_ack:
-                response_queue = await ProxyQueue.create_for_identifier(message["identifier"])
+                response_queue = await ProxyQueue.create_for_identifier(
+                    message["identifier"]
+                )
                 yield await response_queue.dequeue()
                 await response_queue.delete()
             else:
@@ -51,7 +60,9 @@ class WebsocketProxy(WebSocketEndpoint):
         while True:
             message: WebsocketMessage = await response_queue.dequeue()
             logger.debug(message)
-            await websocket.send_text(b64decode(message['payload']['body'].encode()).decode())
+            await websocket.send_text(
+                b64decode(message["payload"]["body"].encode()).decode()
+            )
 
     async def on_connect(self, websocket: WebSocket) -> None:
         message = WebsocketMessage(
@@ -60,18 +71,21 @@ class WebsocketProxy(WebSocketEndpoint):
             payload=WebsocketConnectData(
                 path=get_path_with_query_string(websocket),
                 headers=[
-                    (k.decode(), v.decode())
-                    for k, v
-                    in websocket.scope['headers']
+                    (k.decode(), v.decode()) for k, v in websocket.scope["headers"]
                 ],
-            )
+            ),
         )
 
         async with self.proxy(websocket, message) as m:
-            if m is not None and WebsocketMessageType(m['type']) == WebsocketMessageType.ack:
+            if (
+                m is not None
+                and WebsocketMessageType(m["type"]) == WebsocketMessageType.ack
+            ):
                 await super().on_connect(websocket)
 
-        self.websocket_listen_task = asyncio.create_task(self.listen_for_messages(websocket))
+        self.websocket_listen_task = asyncio.create_task(
+            self.listen_for_messages(websocket)
+        )
 
         def callback(*args, **kwargs):
             self.websocket_listen_task = None
@@ -92,7 +106,7 @@ class WebsocketProxy(WebSocketEndpoint):
             identifier=self.id,
             payload=WebsocketMessageData(
                 body=b64encode(data_bytes).decode(),
-            )
+            ),
         )
 
         async with self.proxy(websocket, message):
@@ -104,7 +118,7 @@ class WebsocketProxy(WebSocketEndpoint):
             identifier=self.id,
             payload=WebsocketDisconnectData(
                 close_code=close_code,
-            )
+            ),
         )
 
         async with self.proxy(websocket, message):