From 72e6b71741a2128098e5a36b191830cc4c350cb4 Mon Sep 17 00:00:00 2001
From: Guilhem Allaman <40383801+gounux@users.noreply.github.com>
Date: Wed, 24 Jul 2024 09:28:33 +0200
Subject: [PATCH] Notify websockets with internal messages containing number of
users in room (#19)
---
README.md | 1 +
gischat/__init__.py | 1 +
gischat/app.py | 34 +++++++++++++++++++++++++++-----
gischat/models.py | 5 +++++
gischat/templates/ws-page.html | 14 ++++++++-----
tests/test_websocket.py | 36 ++++++++++++++++++++++++++++++++--
6 files changed, 79 insertions(+), 12 deletions(-)
diff --git a/README.md b/README.md
index 10dffb7..648a7d3 100644
--- a/README.md
+++ b/README.md
@@ -41,6 +41,7 @@ Following instances are up and running :
- Number of connected users can be fetched using [the `/status` endpoint](https://gischat.geotribu.net/docs#/default/get_status_status_get)
- New users must connect a websocket to the `/room/{room_name}/ws` endpoint
- Messages passing through the websocket are simple JSON dicts like this: `{"message": "hello", "author": "Hans Hibbel"}`
+- :warning: Messages having the `"internal"` author are internal messages and should not be printed, they contain technical information: `{"author": "internal", "nb_users": 36}`
## Deploy a self-hosted instance
diff --git a/gischat/__init__.py b/gischat/__init__.py
index e69de29..95cafed 100755
--- a/gischat/__init__.py
+++ b/gischat/__init__.py
@@ -0,0 +1 @@
+INTERNAL_MESSAGE_AUTHOR = "internal"
diff --git a/gischat/app.py b/gischat/app.py
index b60acb1..2a369a0 100755
--- a/gischat/app.py
+++ b/gischat/app.py
@@ -10,7 +10,14 @@
from fastapi.templating import Jinja2Templates
from starlette.websockets import WebSocketDisconnect
-from gischat.models import MessageModel, RulesModel, StatusModel, VersionModel
+from gischat import INTERNAL_MESSAGE_AUTHOR
+from gischat.models import (
+ InternalMessageModel,
+ MessageModel,
+ RulesModel,
+ StatusModel,
+ VersionModel,
+)
from gischat.utils import get_poetry_version
# logger
@@ -34,6 +41,8 @@ class WebsocketNotifier:
Class used to broadcast messages to registered websockets
"""
+ connections: dict[str, list[WebSocket]]
+
def __init__(self):
# registered websockets for rooms
self.connections: dict[str, list[WebSocket]] = {
@@ -46,17 +55,17 @@ async def get_notification_generator(self):
room, message = yield
await self.notify(room, message)
- async def push(self, msg: str):
+ async def push(self, msg: str) -> None:
await self.generator.asend(msg)
- async def connect(self, room: str, websocket: WebSocket):
+ async def connect(self, room: str, websocket: WebSocket) -> None:
await websocket.accept()
self.connections[room].append(websocket)
- def remove(self, room: str, websocket: WebSocket):
+ def remove(self, room: str, websocket: WebSocket) -> None:
self.connections[room].remove(websocket)
- async def notify(self, room: str, message: str):
+ async def notify(self, room: str, message: str) -> None:
living_connections = []
while len(self.connections[room]) > 0:
# Looping like this is necessary in case a disconnection is handled
@@ -66,6 +75,15 @@ async def notify(self, room: str, message: str):
living_connections.append(websocket)
self.connections[room] = living_connections
+ def get_nb_users(self, room: str) -> int:
+ return len(self.connections[room])
+
+ async def notify_internal(self, room: str) -> None:
+ message = InternalMessageModel(
+ author=INTERNAL_MESSAGE_AUTHOR, nb_users=self.get_nb_users(room)
+ )
+ await self.notify(room, json.dumps(jsonable_encoder(message)))
+
notifier = WebsocketNotifier()
@@ -122,10 +140,15 @@ async def put_message(room: str, message: MessageModel) -> MessageModel:
@app.websocket("/room/{room}/ws")
async def websocket_endpoint(websocket: WebSocket, room: str) -> None:
+
+ # check if room is registered
if room not in notifier.connections.keys():
raise HTTPException(status_code=404, detail=f"Room '{room}' not registered")
+
await notifier.connect(room, websocket)
+ await notifier.notify_internal(room)
logger.info(f"New websocket connected in room '{room}'")
+
try:
while True:
data = await websocket.receive_text()
@@ -134,4 +157,5 @@ async def websocket_endpoint(websocket: WebSocket, room: str) -> None:
await notifier.notify(room, json.dumps(jsonable_encoder(message)))
except WebSocketDisconnect:
notifier.remove(room, websocket)
+ await notifier.notify_internal(room)
logger.info(f"Websocket disconnected from room '{room}'")
diff --git a/gischat/models.py b/gischat/models.py
index d8af7f9..6e17b9b 100755
--- a/gischat/models.py
+++ b/gischat/models.py
@@ -26,3 +26,8 @@ class MessageModel(BaseModel):
def __str__(self) -> str:
return f"[{self.author}]: '{self.message}'"
+
+
+class InternalMessageModel(BaseModel):
+ author: str
+ nb_users: int
diff --git a/gischat/templates/ws-page.html b/gischat/templates/ws-page.html
index c87cbf6..ab51471 100644
--- a/gischat/templates/ws-page.html
+++ b/gischat/templates/ws-page.html
@@ -5,19 +5,20 @@
gischat websocket client
-
@@ -85,7 +86,10 @@ gischat websocket client
}
websocket.onmessage = (event) => {
const data = JSON.parse(event.data);
- const log = `[${data.author}] (${new Date().toLocaleTimeString()}): ${data.message}`;
+ const author = data.author;
+ const log = author === "internal" ?
+ `${data.nb_users} user${data.nb_users > 1 ? "s" : ""} connected in room` :
+ `[${author}] (${new Date().toLocaleTimeString()}): ${data.message}`;
displayMessage(log);
};
websocket.onerror = (error) => {
diff --git a/tests/test_websocket.py b/tests/test_websocket.py
index 5b6f953..c40f4ba 100644
--- a/tests/test_websocket.py
+++ b/tests/test_websocket.py
@@ -3,14 +3,26 @@
import pytest
from fastapi.testclient import TestClient
+from gischat import INTERNAL_MESSAGE_AUTHOR
from tests.conftest import test_rooms
TEST_MESSAGE = "Is this websocket working ?"
+@pytest.mark.parametrize("room", test_rooms())
+def test_websocket_connection(client: TestClient, room: str):
+ with client.websocket_connect(f"/room/{room}/ws") as websocket:
+ data = websocket.receive_json()
+ assert data == {"author": INTERNAL_MESSAGE_AUTHOR, "nb_users": 1}
+
+
@pytest.mark.parametrize("room", test_rooms())
def test_websocket_put_message(client: TestClient, room: str):
with client.websocket_connect(f"/room/{room}/ws") as websocket:
+ assert websocket.receive_json() == {
+ "author": INTERNAL_MESSAGE_AUTHOR,
+ "nb_users": 1,
+ }
client.put(
f"/room/{room}/message",
json={"message": TEST_MESSAGE, "author": f"ws-tester-{room}"},
@@ -22,6 +34,10 @@ def test_websocket_put_message(client: TestClient, room: str):
@pytest.mark.parametrize("room", test_rooms())
def test_websocket_send_message(client: TestClient, room: str):
with client.websocket_connect(f"/room/{room}/ws") as websocket:
+ assert websocket.receive_json() == {
+ "author": INTERNAL_MESSAGE_AUTHOR,
+ "nb_users": 1,
+ }
websocket.send_json({"message": TEST_MESSAGE, "author": f"ws-tester-{room}"})
data = websocket.receive_json()
assert data == {"message": TEST_MESSAGE, "author": f"ws-tester-{room}"}
@@ -38,9 +54,25 @@ def nb_connected_users(json: dict[str, Any], room: str) -> bool:
@pytest.mark.parametrize("room", test_rooms())
def test_websocket_nb_users_connected(client: TestClient, room: str):
assert nb_connected_users(client.get("/status").json(), room) == 0
- with client.websocket_connect(f"/room/{room}/ws"):
+ with client.websocket_connect(f"/room/{room}/ws") as websocket1:
+ assert websocket1.receive_json() == {
+ "author": INTERNAL_MESSAGE_AUTHOR,
+ "nb_users": 1,
+ }
assert nb_connected_users(client.get("/status").json(), room) == 1
- with client.websocket_connect(f"/room/{room}/ws"):
+ with client.websocket_connect(f"/room/{room}/ws") as websocket2:
+ assert websocket1.receive_json() == {
+ "author": INTERNAL_MESSAGE_AUTHOR,
+ "nb_users": 2,
+ }
+ assert websocket2.receive_json() == {
+ "author": INTERNAL_MESSAGE_AUTHOR,
+ "nb_users": 2,
+ }
assert nb_connected_users(client.get("/status").json(), room) == 2
+ assert websocket1.receive_json() == {
+ "author": INTERNAL_MESSAGE_AUTHOR,
+ "nb_users": 1,
+ }
assert nb_connected_users(client.get("/status").json(), room) == 1
assert nb_connected_users(client.get("/status").json(), room) == 0