Add API to disconnect clients

Author: Pithikos

README.md

@@ -79,8 +79,10 @@ The WebsocketServer can be initialized with the below parameters.
 | `set_fn_message_received()` | Sets a callback function that will be called when a `client` sends a message          | function        | None  |
 | `send_message()`            | Sends a `message` to a specific `client`. The message is a simple string.             | client, message | None  |
 | `send_message_to_all()`     | Sends a `message` to **all** connected clients. The message is a simple string.       | message         | None  |
-| `shutdown_gracefully()`     | Shutdown server by sending a websocket CLOSE handshake to all connected clients.      | None            | None  |
-| `shutdown_abruptly()`       | Shutdown server without sending any websocket CLOSE handshake.                        | None            | None  |
+| `disconnect_clients_gracefully()` | Disconnect all connected clients by sending a websocket CLOSE handshake.        | Optional: status, reason | None  |
+| `disconnect_clients_abruptly()`   | Disconnect all connected clients. Clients won't be aware until they try to send some data. | None | None  |
+| `shutdown_gracefully()`     | Disconnect clients with a CLOSE handshake and shutdown server.      | Optional: status, reason | None  |
+| `shutdown_abruptly()`       | Disconnect clients and shutdown server with no handshake.      | None            | None  |
 
 
 

tests/test_server.py

@@ -115,3 +115,29 @@ def test_client_closes_gracefully(self, session):
         assert not server.clients
         with pytest.raises(BrokenPipeError):
             old_client_handler.connection.send(b"test")
+
+    def test_disconnect_clients_abruptly(self, session):
+        client, server = session
+        assert client.connected
+        assert server.clients
+        server.disconnect_clients_abruptly()
+        assert not server.clients
+
+        # Client won't be aware until trying to write more data
+        with pytest.raises(BrokenPipeError):
+            for i in range(3):
+                client.send("test")
+                sleep(0.2)
+
+    def test_disconnect_clients_gracefully(self, session):
+        client, server = session
+        assert client.connected
+        assert server.clients
+        server.disconnect_clients_gracefully()
+        assert not server.clients
+
+        # Client won't be aware until trying to write more data
+        with pytest.raises(BrokenPipeError):
+            for i in range(3):
+                client.send("test")
+                sleep(0.2)

websocket_server/websocket_server.py

@@ -86,6 +86,12 @@ def shutdown_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_R
     def shutdown_abruptly(self):
         self._shutdown_abruptly()
 
+    def disconnect_clients_gracefully(self):
+        self._disconnect_clients_gracefully()
+
+    def disconnect_clients_abruptly(self):
+        self._disconnect_clients_abruptly()
+
 
 class WebsocketServer(ThreadingMixIn, TCPServer, API):
     """
@@ -196,12 +202,7 @@ def _shutdown_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_
         Send a CLOSE handshake to all connected clients before terminating server
         """
         self.keep_alive = False
-
-        # Send CLOSE to clients
-        for client in self.clients:
-            client["handler"].send_close(CLOSE_STATUS_NORMAL, reason)
-
-        self._terminate_client_handlers()
+        self._disconnect_clients_gracefully(status, reason)
         self.server_close()
         self.shutdown()
 
@@ -210,10 +211,24 @@ def _shutdown_abruptly(self):
         Terminate server without sending a CLOSE handshake
         """
         self.keep_alive = False
-        self._terminate_client_handlers()
+        self._disconnect_clients_abruptly()
         self.server_close()
         self.shutdown()
 
+    def _disconnect_clients_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
+        """
+        Terminate clients gracefully without shutting down the server
+        """
+        for client in self.clients:
+            client["handler"].send_close(CLOSE_STATUS_NORMAL, reason)
+        self._terminate_client_handlers()
+
+    def _disconnect_clients_abruptly(self):
+        """
+        Terminate clients abruptly (no CLOSE handshake) without shutting down the server
+        """
+        self._terminate_client_handlers()
+
 
 class WebSocketHandler(StreamRequestHandler):