add web socket server with Fastapi

jbcodeforce · Apr 22, 2024 · a652359 · a652359
1 parent 7a12aa8
commit a652359
Show file tree

Hide file tree

Showing 4 changed files with 129 additions and 1 deletion.
diff --git a/docs/python/faq.md b/docs/python/faq.md
@@ -275,4 +275,8 @@ A scope is a textual region of a Python program, where a namespace is directly a
     new_labels = ['${}'.format(int(amt)) for amt in ticks]
     # Set the new labels
     graph.set_yticklabels(new_labels)
-```
+```
+
+
+???- question "How to program web socket server"
+    See [FastAPI doc](https://fastapi.tiangolo.com/advanced/websockets/) with [testing](https://fastapi.tiangolo.com/advanced/testing-websockets/) and matching code in [](https://github.com/jbcodeforce/python-code/tree/master/web-server/websocket_server)
diff --git a/web-server/websocket_server/TestWSServer.py b/web-server/websocket_server/TestWSServer.py
@@ -0,0 +1,16 @@
+from fastapi.testclient import TestClient
+from main import app
+
+def test_read_main():
+    client = TestClient(app)
+    response = client.get("/hello")
+    assert response.status_code == 200
+    assert response.json() == {"msg": "Hello World"}
+
+
+def test_websocket():
+    client = TestClient(app)
+    with client.websocket_connect("/ws/12") as websocket:
+        websocket.send_text("Hello from client 12")
+        data = websocket.receive_text()
+        assert data == "You wrote: Hello from client 12"
diff --git a/web-server/websocket_server/a_client.py b/web-server/websocket_server/a_client.py
@@ -0,0 +1,17 @@
+import websocket
+
+# Set the WebSocket endpoint URL
+url = "ws://localhost:8001/ws"
+
+# Create a WebSocket connection
+ws = websocket.create_connection(url)
+
+# Send a message to the server
+ws.send("Hello, server!")
+
+# Receive a response from the server
+response = ws.recv()
+print(f"Received: {response}")
+
+# Close the WebSocket connection
+ws.close()
diff --git a/web-server/websocket_server/main.py b/web-server/websocket_server/main.py
@@ -0,0 +1,91 @@
+from fastapi import FastAPI, WebSocket,WebSocketDisconnect 
+from fastapi.responses import HTMLResponse
+
+"""
+ simple example to demonstrate how to handle and broadcast messages to several WebSocket connections.
+"""
+
+BACKEND_URL="localhost:8001"
+
+html = f"""
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>Chat</title>
+    </head>
+    <body>
+        <h1>WebSocket Chat</h1>
+        <h2>Your ID: <span id="ws-id"></span></h2>
+        <form action="" onsubmit="sendMessage(event)">
+            <input type="text" id="messageText" autocomplete="off"/>
+            <button>Send</button>
+        </form>
+        <ul id='messages'>
+        </ul>
+        <script>
+            var client_id = Date.now()
+            document.querySelector("#ws-id").textContent = client_id;
+            var ws = new WebSocket(`ws://{BACKEND_URL}/ws/${{client_id}}`);
+            ws.onmessage = function(event) {{
+                var messages = document.getElementById('messages')
+                var message = document.createElement('li')
+                var content = document.createTextNode(event.data)
+                message.appendChild(content)
+                messages.appendChild(message)
+            }};
+            function sendMessage(event) {{
+                var input = document.getElementById("messageText")
+                ws.send(input.value)
+                input.value = ''
+                event.preventDefault()
+            }}
+        </script>
+    </body>
+</html>
+"""
+
+app = FastAPI()
+
+class ConnectionManager:
+    """
+    Used to keep active connection with the multiple clients
+    """
+    def __init__(self):
+        self.active_connections: list[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+
+    async def send_personal_message(self, message: str, websocket: WebSocket):
+        await websocket.send_text(message)
+
+    async def broadcast(self, message: str):
+        for connection in self.active_connections:
+            await connection.send_text(message)
+
+
+manager = ConnectionManager()
+
+@app.get("/")
+async def get():
+    return HTMLResponse(html)
+
+@app.get("/hello")
+async def get_hello():
+    return { "msg" : "Hello World"}
+
+@app.websocket("/ws/{client_id}")
+async def websocket_endpoint(websocket: WebSocket, client_id: str):
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive_text()
+            await manager.send_personal_message(f"You wrote: {data}", websocket)
+            await manager.broadcast(f"Client #{client_id} says: {data}")
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        await manager.broadcast(f"Client #{client_id} left the chat")