From 69531e464a3ed56519f8f4e31bc27abdd4f4c576 Mon Sep 17 00:00:00 2001
From: Vukasin Gostovic <vukasin@gostovic.me>
Date: Sun, 7 Jan 2024 23:02:45 +0100
Subject: [PATCH] ws logic

---
 hrotti/main.py              | 12 ++++--------
 hrotti/methods.py           |  2 +-
 hrotti/websocket_manager.py | 14 ++++++++++++++
 3 files changed, 19 insertions(+), 9 deletions(-)
 create mode 100644 hrotti/websocket_manager.py

diff --git a/hrotti/main.py b/hrotti/main.py
index c0fd887..c600574 100644
--- a/hrotti/main.py
+++ b/hrotti/main.py
@@ -2,7 +2,8 @@
 
 from fastapi import FastAPI, WebSocket
 
-from .methods import RPCRequest, BlockInfo, handle_request
+from .methods import BlockInfo, RPCRequest, handle_request
+from .websocket_manager import accept_ws_request
 
 app = FastAPI()
 
@@ -17,16 +18,11 @@
 
 @app.post("/http")
 async def handle_json_rpc(request: RPCRequest):
-    return handle_request(request, info)
+    return handle_request(info, request)
 
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     while True:
-        data = await websocket.receive_text()
-        try:
-            request = json.loads(data)
-            await websocket.send_text(handle_request(request.get("method"), info))
-        except json.JSONDecodeError:
-            await websocket.send_text("Error: Invalid JSON")
+        accept_ws_request(info, websocket)
diff --git a/hrotti/methods.py b/hrotti/methods.py
index fc0001b..03a3261 100644
--- a/hrotti/methods.py
+++ b/hrotti/methods.py
@@ -72,7 +72,7 @@ def return_block(head, transaction):
 }
 
 
-def handle_request(request: RPCRequest, info: BlockInfo):
+def handle_request(info: BlockInfo, request: RPCRequest):
     try:
         print(request.method)
         match request.method:
diff --git a/hrotti/websocket_manager.py b/hrotti/websocket_manager.py
new file mode 100644
index 0000000..dc21a7b
--- /dev/null
+++ b/hrotti/websocket_manager.py
@@ -0,0 +1,14 @@
+import json
+
+from fastapi import WebSocket
+
+from .methods import BlockInfo, handle_request
+
+
+async def accept_ws_request(info: BlockInfo, websocket: WebSocket):
+    data = await websocket.receive_text()
+    try:
+        request = json.loads(data)
+        await websocket.send_text(handle_request(request.get("method"), info))
+    except json.JSONDecodeError:
+        await websocket.send_text("Error: Invalid JSON")