From 67b9db2ed0d6ae43ae7e9e0a8d061e068ac41b46 Mon Sep 17 00:00:00 2001
From: Eden Kirin <eden.kirin@intis.hr>
Date: Fri, 24 Mar 2023 11:39:36 +0100
Subject: [PATCH] Cleanups and optimizations

---
 README.md               |  1 +
 src/frontend/index.html |  5 ++++-
 src/funnel/Makefile     |  4 ++++
 src/funnel/main.py      | 30 ++++++++----------------------
 src/funnel/ws_server.py | 13 ++++++-------
 5 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index 3bcac97..f401eb0 100644
--- a/README.md
+++ b/README.md
@@ -61,6 +61,7 @@ deactivate FE
 | Funnel | Python | - |
 | ServeCurrentTime | NodeJS | 50000 |
 | ServeSegments | Go | 50001 |
+| WS Server | Go | 50010 |
 
 ### Funnel
 
diff --git a/src/frontend/index.html b/src/frontend/index.html
index 4bc2ddd..185f451 100644
--- a/src/frontend/index.html
+++ b/src/frontend/index.html
@@ -16,8 +16,11 @@
 </body>
 
 <script type="module">
+    const WS_HOST = "localhost";
+    const WS_PORT = 50010;
+
     window.onload = function () {
-        const ws = new WebSocket('ws://localhost:5000');
+        const ws = new WebSocket(`ws://${WS_HOST}:${WS_PORT}`);
 
         ws.onmessage = function (event) {
             const data = JSON.parse(event.data);
diff --git a/src/funnel/Makefile b/src/funnel/Makefile
index d06e837..ff3cb81 100644
--- a/src/funnel/Makefile
+++ b/src/funnel/Makefile
@@ -5,6 +5,10 @@ run:
 	@source env/bin/activate && \
 		python main.py
 
+run-ws:
+	@source env/bin/activate && \
+		python ws_server.py
+
 proto:
 	@source env/bin/activate && \
 		python \
diff --git a/src/funnel/main.py b/src/funnel/main.py
index f875b4e..7faffa4 100644
--- a/src/funnel/main.py
+++ b/src/funnel/main.py
@@ -2,11 +2,11 @@ import asyncio
 import json
 import time
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Optional
+from typing import AsyncGenerator
 
 import grpc
 import websockets
-from websockets.server import WebSocketServerProtocol
+from websockets.legacy.client import connect as ws_connect
 
 from stubs.serve_segments_pb2 import (
     GetHoursRequest,
@@ -25,7 +25,7 @@ SERVE_CURRENTTIME_PORT = 50000
 SERVE_SEGMENTS_HOST = "localhost"
 SERVE_SEGMENTS_PORT = 50001
 WS_HOST = "localhost"
-WS_PORT = 5000
+WS_PORT = 50010
 
 TIMEZONE = "Europe/Zagreb"
 
@@ -39,26 +39,26 @@ async def get_serve_segments_stub() -> AsyncGenerator[ServeSegmentsStub, None]:
         yield ServeSegmentsStub(channel)
 
 
-async def get_hours(stub: ServeSegmentsStub) -> Optional[int]:
+async def get_hours(stub: ServeSegmentsStub) -> int:
     response: GetHoursResponse = await stub.GetHours(GetHoursRequest(timezone=TIMEZONE))
     return response.hours
 
 
-async def get_minutes(stub: ServeSegmentsStub) -> Optional[int]:
+async def get_minutes(stub: ServeSegmentsStub) -> int:
     response: GetMinutesResponse = await stub.GetMinutes(
         GetMinutesRequest(timezone=TIMEZONE)
     )
     return response.minutes
 
 
-async def get_seconds(stub: ServeSegmentsStub) -> Optional[int]:
+async def get_seconds(stub: ServeSegmentsStub) -> int:
     response: GetSecondsResponse = await stub.GetSeconds(
         GetSecondsRequest(timezone=TIMEZONE)
     )
     return response.seconds
 
 
-async def get_milliseconds(stub: ServeSegmentsStub) -> Optional[int]:
+async def get_milliseconds(stub: ServeSegmentsStub) -> int:
     response: GetMillisecondsResponse = await stub.GetMilliseconds(
         GetMillisecondsRequest(timezone=TIMEZONE)
     )
@@ -67,7 +67,7 @@ async def get_milliseconds(stub: ServeSegmentsStub) -> Optional[int]:
 
 async def main():
     ws_uri = f"ws://{WS_HOST}:{WS_PORT}"
-    async with websockets.connect(uri=ws_uri) as websocket:
+    async with ws_connect(uri=ws_uri) as websocket:
         async with get_serve_segments_stub() as stub:
             while True:
                 t = time.perf_counter()
@@ -106,22 +106,8 @@ async def main():
                 await websocket.send(json.dumps(data))
                 await websocket.recv()
 
-                # time.sleep(0.5)
-
                 print(f"RESULT: {formatted_time}, T: {comm_cycle}")
 
 
-async def send_ws_message():
-    ws_uri = f"ws://{WS_HOST}:{WS_PORT}"
-    async with websockets.connect(ws_uri) as websocket:
-        data = {
-            "message": "some text",
-            "value": 12345,
-        }
-
-        await websocket.send(json.dumps(data))
-
-
 if __name__ == "__main__":
     asyncio.run(main())
-    # asyncio.run(send_ws_message())
diff --git a/src/funnel/ws_server.py b/src/funnel/ws_server.py
index 589449e..e587f59 100644
--- a/src/funnel/ws_server.py
+++ b/src/funnel/ws_server.py
@@ -1,18 +1,17 @@
 import asyncio
-import json
-import websockets
-from websockets import broadcast
+from websockets.legacy.protocol import broadcast
+from websockets.legacy.server import serve as ws_serve
 from websockets.server import WebSocketServerProtocol
 
 WS_HOST = "localhost"
-WS_PORT = 5000
+WS_PORT = 50010
 
 connected_clients = set()
 
 
 async def ws_handler(websocket: WebSocketServerProtocol):
     connected_clients.add(websocket)
-    print("Client add: ", websocket)
+    print("Client added: ", websocket)
 
     try:
         async for message in websocket:
@@ -20,12 +19,12 @@ async def ws_handler(websocket: WebSocketServerProtocol):
             # await websocket.send(f"Are you talking to me? {message}")
     finally:
         connected_clients.remove(websocket)
-        print("Client remove: ", websocket)
+        print("Client removed: ", websocket)
 
 
 async def main():
     print(f"Starting WS server on {WS_HOST}:{WS_PORT}")
-    async with websockets.serve(
+    async with ws_serve(
         ws_handler=ws_handler,
         host=WS_HOST,
         port=WS_PORT,