Implement shutdown_gracefully

Author: Pithikos

tests/test_server.py

@@ -5,7 +5,9 @@
 
 
 def test_send_close(client_session):
-    "Ensure client stops receiving data once we send_close (socket is still open)"
+    """
+    Ensure client stops receiving data once we send_close (socket is still open)
+    """
     client, server = client_session
     assert client.received_messages == []
 
@@ -21,6 +23,20 @@ def test_send_close(client_session):
     assert client.received_messages == ["test1"]
 
 
+def test_shutdown_gracefully(client_session):
+    client, server = client_session
+    assert client.ws.sock and client.ws.sock.connected
+    assert server.socket.fileno() > 0
+
+    server.shutdown_gracefully()
+    sleep(0.5)
+
+    # Ensure all parties disconnected
+    assert not client.ws.sock
+    assert server.socket.fileno() == -1
+    assert not server.clients
+
+
 def test_client_closes_gracefully(session):
     client, server = session
     assert client.connected

websocket_server/websocket_server.py

@@ -86,6 +86,18 @@ def send_message(self, client, msg):
     def send_message_to_all(self, msg):
         self._multicast(msg)
 
+    def shutdown_gracefully(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
+        """
+        Close with a websocket handshake
+
+        1. Send CLOSE to all clients
+        2. Close TCP
+        """
+        self.keep_alive = False
+        for client in self.clients:
+            client["handler"].send_close(CLOSE_STATUS_NORMAL, reason)
+        self.server_close()
+
 
 class WebsocketServer(ThreadingMixIn, TCPServer, API):
     """
@@ -258,7 +270,16 @@ def send_close(self, status=CLOSE_STATUS_NORMAL, reason=DEFAULT_CLOSE_REASON):
         """
         if status < CLOSE_STATUS_NORMAL or status > 1015:
             raise Exception(f"CLOSE status must be between 1000 and 1015, got {status}")
-        self.request.send(struct.pack('!H', status) + reason, OPCODE_CLOSE_CONN)
+
+        header = bytearray()
+        payload = struct.pack('!H', status) + reason
+        payload_length = len(payload)
+        assert payload_length <= 125, "We only support short closing reasons at the moment"
+
+        # Send CLOSE with status & reason
+        header.append(FIN | OPCODE_CLOSE_CONN)
+        header.append(payload_length)
+        self.request.send(header + payload)
 
     def send_text(self, message, opcode=OPCODE_TEXT):
         """