➕ Add heartbeat handling of reply.

Author: DanielAtCosmicDNA

plover_websocket_server/websocket/server.py

@@ -120,6 +120,11 @@ async def websocket_handler(self, request: Request) -> WebSocketResponse:
         try:
             async for message in socket:
                 if message.type == WSMsgType.TEXT:
+                    # Handle unencrypted heartbeat first
+                    if message.data == "ping":
+                        await socket.send_str("pong")
+                        continue
+
                     try:
                         log.debug("Decrypting message...")
                         decrypted: dict = mail_box.unbox(message.data)

pyproject.toml

@@ -49,7 +49,8 @@ include-package-data = false
 [project.optional-dependencies]
 test = [
     "pytest",
-    "PyQt5"
+    "PyQt5",
+    "PySide6"
 ]
 
 dev = [

tests/client_example.py

@@ -181,6 +181,13 @@ async def send_websocket_message(self, message) -> None:
         """
         await self.socket.send_str(self.encrypt(message))
 
+    async def ping(self) -> None:
+        """
+        Sends a raw 'ping' message to the server for heartbeat testing.
+
+        """
+        await self.socket.send_str("ping")
+
     async def connect_to_websocket(self, message) -> None:
         """
         Connects to the server using the WebSocket protocol.
@@ -195,12 +202,21 @@ async def connect_to_websocket(self, message) -> None:
             url,
             params=self._get_encryption_params(message),
             headers=self.headers,
-            ssl=self.ssl_context,
+            ssl= False if self.ssl_context is None else True,
         )
 
     async def receive_json(self):
         return await self.socket.receive_json()
 
+    async def receive_raw(self):
+        """
+        Receives a raw message from the websocket.
+
+        Returns:
+            any: The raw data from the message.
+        """
+        return (await self.socket.receive()).data
+
     async def disconnect(self) -> None:
         """
         Disconnects the client from the server by closing the session.

tests/test_server.py

@@ -25,6 +25,21 @@ async def websocket(client: Client):
     assert reply == message
 
 
+async def ping_pong(client: Client):
+    """
+    Callback function for testing websocket heartbeat (ping-pong).
+
+    Args:
+        client (Client): The client object representing the websocket connection.
+
+    Returns:
+        None
+    """
+    await client.connect_to_websocket({"message": "ping-pong test"})
+    await client.ping()
+    reply = await client.receive_raw()
+    assert reply == "pong"
+
 async def http(client: Client):
     """
     This function is the callback for HTTP requests.
@@ -38,7 +53,7 @@ async def http(client: Client):
     pass
 
 
-@mark.parametrize("callback", [http, websocket])
+@mark.parametrize("callback", [http, websocket, ping_pong])
 def test_server(callback: Awaitable) -> None:
     """
     Test the server functionality.