Skip to content

Commit

Permalink
Notify websockets with internal messages containing number of users i…
Browse files Browse the repository at this point in the history
…n room (#19)
  • Loading branch information
gounux authored Jul 24, 2024
1 parent 0ba51d5 commit 72e6b71
Show file tree
Hide file tree
Showing 6 changed files with 79 additions and 12 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
1 change: 1 addition & 0 deletions gischat/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
INTERNAL_MESSAGE_AUTHOR = "internal"
34 changes: 29 additions & 5 deletions gischat/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]] = {
Expand All @@ -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
Expand All @@ -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()

Expand Down Expand Up @@ -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()
Expand All @@ -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}'")
5 changes: 5 additions & 0 deletions gischat/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
14 changes: 9 additions & 5 deletions gischat/templates/ws-page.html
Original file line number Diff line number Diff line change
Expand Up @@ -5,19 +5,20 @@
</head>
<body>
<h1>gischat websocket client</h1>
<form action="" onsubmit="sendMessage(event)">
<form>
<label>Instance: <input type="text" id="instance" autocomplete="off" value=""/></label>
<label>SSL: <input type="checkbox" id="ssl" checked/></label>
<button onclick="onRulesButtonClick(event)">Rules</button>
<button onclick="onStatusButtonClick(event)">Status</button>
<hr>
<label>Room: <input type="text" id="roomId" autocomplete="off" value=""/></label>
<label>Room: <input type="text" id="roomId" autocomplete="off" value="QGIS"/></label>
<button onclick="onDetectButtonClick(event)">Detect</button>
<button id="connectButton" onclick="onConnectButtonClick(event)">Connect</button>
<hr>
<br>
<label>Author: <input type="text" id="authorId" autocomplete="off" value=""/></label>
<hr>
<label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
<button>Send</button>
<button onclick="sendMessage(event)">Send</button>
</form>
<hr>
<ul id='messages'>
Expand Down Expand Up @@ -85,7 +86,10 @@ <h1>gischat websocket client</h1>
}
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) => {
Expand Down
36 changes: 34 additions & 2 deletions tests/test_websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"},
Expand All @@ -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}"}
Expand All @@ -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

0 comments on commit 72e6b71

Please sign in to comment.