Fixed pre-commit issuesv3

5 files changed, 78 insertions, 70 deletions

diff --git a/proxy/__init__.py b/proxy/__init__.py
index fcdb23b..aba8ad2 100644
--- a/proxy/__init__.py
+++ b/proxy/__init__.py
@@ -7,7 +7,7 @@ from proxy.websockets import WebsocketProxy
 app = Starlette(
     debug=True,
     routes=[
-        Route('/{path:path}', Proxy),
-        WebSocketRoute('/{path:path}', WebsocketProxy),
-    ]
+        Route("/{path:path}", Proxy),
+        WebSocketRoute("/{path:path}", WebsocketProxy),
+    ],
 )
diff --git a/proxy/endpoints.py b/proxy/endpoints.py
index 0a23e7b..f514226 100644
--- a/proxy/endpoints.py
+++ b/proxy/endpoints.py
@@ -1,13 +1,13 @@
 import logging
-from base64 import b64encode, b64decode
+from base64 import b64decode, b64encode
 from uuid import uuid4
 
 from starlette.endpoints import HTTPEndpoint
 from starlette.requests import Request
 from starlette.responses import Response
 
-from proxy.utils import get_path_with_query_string
 from proxy.queue import ProxyQueue
+from proxy.utils import get_path_with_query_string
 from ttun_server.types import HttpMessage, HttpMessageType, HttpRequestData
 
 logger = logging.getLogger(__name__)
@@ -18,16 +18,15 @@ class HeaderMapping:
         self._headers = headers
 
     def items(self):
-        for header in self._headers:
-            yield header
+        yield from self._headers
 
 
 class Proxy(HTTPEndpoint):
     async def dispatch(self) -> None:
         request = Request(self.scope, self.receive)
 
-        subdomain = request.path_params['subdomain']
-        response = Response(content='Not Found', status_code=404)
+        subdomain = request.path_params["subdomain"]
+        response = Response(content="Not Found", status_code=404)
 
         identifier = str(uuid4())
         response_queue = await ProxyQueue.create_for_identifier(identifier)
@@ -35,7 +34,7 @@ class Proxy(HTTPEndpoint):
         try:
             request_queue = await ProxyQueue.get_for_identifier(subdomain)
 
-            logger.debug('PROXY %s%s ', subdomain, request.url)
+            logger.debug("PROXY %s%s ", subdomain, request.url)
             await request_queue.enqueue(
                 HttpMessage(
                     type=HttpMessageType.request.value,
@@ -44,17 +43,17 @@ class Proxy(HTTPEndpoint):
                         method=request.method,
                         path=get_path_with_query_string(request),
                         headers=list(request.headers.items()),
-                        body=b64encode(await request.body()).decode()
-                    )
+                        body=b64encode(await request.body()).decode(),
+                    ),
                 )
             )
 
             _response = await response_queue.dequeue()
-            payload = _response['payload']
+            payload = _response["payload"]
             response = Response(
-                status_code=payload['status'],
-                headers=HeaderMapping(payload['headers']),
-                content=b64decode(payload['body'].encode())
+                status_code=payload["status"],
+                headers=HeaderMapping(payload["headers"]),
+                content=b64decode(payload["body"].encode()),
             )
         except AssertionError:
             pass
diff --git a/proxy/queue.py b/proxy/queue.py
index cfa0f3c..70fe3d0 100644
--- a/proxy/queue.py
+++ b/proxy/queue.py
@@ -2,8 +2,7 @@ import asyncio
 import json
 import logging
 import os
-import traceback
-from typing import Type
+from typing import ClassVar, Self
 
 from ttun_server.redis import RedisConnectionPool
 from ttun_server.types import Message
@@ -16,37 +15,37 @@ class BaseProxyQueue:
         self.identifier = identifier
 
     @classmethod
-    async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue':
-        raise NotImplementedError(f'Please implement create_for_identifier')
+    async def create_for_identifier(cls, identifier: str) -> BaseProxyQueue:
+        raise NotImplementedError("Please implement create_for_identifier")
 
     @classmethod
-    async def get_for_identifier(cls, identifier: str) -> Type['self']:
+    async def get_for_identifier(cls, identifier: str) -> Self:
         assert await cls.has_connection(identifier)
         return cls(identifier)
 
     @classmethod
     async def has_connection(cls, identifier) -> bool:
-        raise NotImplementedError(f'Please implement has_connection')
+        raise NotImplementedError("Please implement has_connection")
 
     async def enqueue(self, message: Message):
-        raise NotImplementedError(f'Please implement send_request')
+        raise NotImplementedError("Please implement send_request")
 
     async def dequeue(self) -> Message:
-        raise NotImplementedError(f'Please implement handle_requests')
+        raise NotImplementedError("Please implement handle_requests")
 
     async def delete(self):
-        raise NotImplementedError(f'Please implement delete')
+        raise NotImplementedError("Please implement delete")
 
 
 class MemoryProxyQueue(BaseProxyQueue):
-    connections: dict[str, asyncio.Queue] = {}
+    connections: ClassVar[dict[str, asyncio.Queue]] = {}
 
     @classmethod
     async def has_connection(cls, identifier) -> bool:
         return identifier in cls.connections
 
     @classmethod
-    async def create_for_identifier(cls, identifier: str) -> 'MemoryProxyQueue':
+    async def create_for_identifier(cls, identifier: str) -> MemoryProxyQueue:
         instance = cls(identifier)
         cls.connections[identifier] = asyncio.Queue()
 
@@ -66,67 +65,59 @@ class RedisProxyQueue(BaseProxyQueue):
     def __init__(self, identifier):
         super().__init__(identifier)
 
-        self.pubsub = RedisConnectionPool()\
-            .get_connection()\
-            .pubsub()
+        self.pubsub = RedisConnectionPool().get_connection().pubsub()
 
         self.subscription_queue = asyncio.Queue()
 
     @classmethod
-    async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue':
+    async def create_for_identifier(cls, identifier: str) -> BaseProxyQueue:
         instance = cls(identifier)
 
-        await instance.pubsub.subscribe(f'request_{identifier}')
+        await instance.pubsub.subscribe(f"request_{identifier}")
         return instance
 
     @classmethod
-    async def get_for_identifier(cls, identifier: str) -> 'RedisProxyQueue':
-        instance: 'RedisProxyQueue' = await super().get_for_identifier(identifier)
+    async def get_for_identifier(cls, identifier: str) -> RedisProxyQueue:
+        instance: RedisProxyQueue = await super().get_for_identifier(identifier)
 
-        await instance.pubsub.subscribe(f'response_{identifier}')
+        await instance.pubsub.subscribe(f"response_{identifier}")
 
         return instance
 
     @classmethod
     async def has_connection(cls, identifier) -> bool:
         logger.debug(await RedisConnectionPool.get_connection().pubsub_channels())
-        return f'request_{identifier}' in {
+        return f"request_{identifier}" in {
             channel.decode()
-            for channel
-            in await RedisConnectionPool \
-                .get_connection() \
-                .pubsub_channels()
+            for channel in await RedisConnectionPool.get_connection().pubsub_channels()
         }
 
     async def wait_for_message(self):
         async for message in self.pubsub.listen():
-            match message['type']:
-                case 'subscribe':
+            match message["type"]:
+                case "subscribe":
                     continue
                 case _:
-                    return message['data']
+                    return message["data"]
 
     async def enqueue(self, message: Message):
-        await RedisConnectionPool \
-            .get_connection() \
-            .publish(f'request_{self.identifier}', json.dumps(message))
+        await RedisConnectionPool.get_connection().publish(
+            f"request_{self.identifier}", json.dumps(message)
+        )
 
     async def dequeue(self) -> Message:
         message = await self.wait_for_message()
         return json.loads(message)
 
     async def delete(self):
-        await self.pubsub.unsubscribe(f'request_{self.identifier}')
+        await self.pubsub.unsubscribe(f"request_{self.identifier}")
 
-        await RedisConnectionPool.get_connection()\
-            .srem('connections', self.identifier)
+        await RedisConnectionPool.get_connection().srem("connections", self.identifier)
 
 
 class ProxyQueueMeta(type):
     def __new__(cls, name, superclasses, attributes):
-        return RedisProxyQueue \
-            if 'REDIS_URL' in os.environ \
-            else MemoryProxyQueue
+        return RedisProxyQueue if "REDIS_URL" in os.environ else MemoryProxyQueue
 
 
 class ProxyQueue(BaseProxyQueue, metaclass=ProxyQueueMeta):
diff --git a/proxy/utils.py b/proxy/utils.py
index 2b80f43..3fff388 100644
--- a/proxy/utils.py
+++ b/proxy/utils.py
@@ -3,5 +3,9 @@ from starlette.requests import HTTPConnection
 
 def get_path_with_query_string(connection: HTTPConnection) -> str:
     path = connection.url.path
-    query_string = '?' + connection.scope['query_string'].decode() if connection.scope['query_string'] else ''
+    query_string = (
+        "?" + connection.scope["query_string"].decode()
+        if connection.scope["query_string"]
+        else ""
+    )
     return f"{path}{query_string}"
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):