Repository: encode/starlette

Issue:
[Feature] Reason for websocket.close
I've recently started actively developing with websockets in starlette. When closing a websocket you can provide a code e.g. 1000, 1008 … to tell the client the basic reason. But the websocket protocol allows the developer to also give UTF-8-encoded data as a complementary description/ reason. This would be useful in multiple cases such as mentioning the specific violitioan for code 1008 (POLICY VIOLIATION). [Websockets](https://websockets.readthedocs.io/en/stable/api.html?highlight=reason#websockets.protocol.WebSocketCommonProtocol.close) supports this too and their [implementation](https://websockets.readthedocs.io/en/stable/_modules/websockets/protocol.html#WebSocketCommonProtocol.close) could maybe help.

Specification: [https://tools.ietf.org/](https://tools.ietf.org/html/rfc6455#section-5.5.1)




Discussion:
Created basic implementation in #992 
https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event

This is not in the ASGI protocol, so it cannot be achieved.
Thanks @abersheeran yup seems correct to me too.

"code" is part of the ASGI spec, but "reason" isn't. If the ASGI spec evolves there at some point, then we could consider it.
["reason" has been added to the ASGI spec.](https://github.com/django/asgiref/pull/235) Will this library add it now as well?
The code from #992 should still be able to be used … Well, up for grabs
Should we wait it to be release on the ASGI website? In any case, PR is more than welcome. Thanks @Chaostheorie :) 
The `reason` attribute is also on the ASGI documentation https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event
When I wrote my comment, it was not :)

They just fixed the docs (I asked them the reason :sweat_smile:). 

Resolving patch:
diff --git a/starlette/testclient.py b/starlette/testclient.py
--- a/starlette/testclient.py
+++ b/starlette/testclient.py
@@ -352,7 +352,9 @@ async def _asgi_send(self, message: Message) -> None:
 
     def _raise_on_close(self, message: Message) -> None:
         if message["type"] == "websocket.close":
-            raise WebSocketDisconnect(message.get("code", 1000))
+            raise WebSocketDisconnect(
+                message.get("code", 1000), message.get("reason", "")
+            )
 
     def send(self, message: Message) -> None:
         self._receive_queue.put(message)
diff --git a/starlette/websockets.py b/starlette/websockets.py
--- a/starlette/websockets.py
+++ b/starlette/websockets.py
@@ -13,8 +13,9 @@ class WebSocketState(enum.Enum):
 
 
 class WebSocketDisconnect(Exception):
-    def __init__(self, code: int = 1000) -> None:
+    def __init__(self, code: int = 1000, reason: str = None) -> None:
         self.code = code
+        self.reason = reason or ""
 
 
 class WebSocket(HTTPConnection):
@@ -146,13 +147,18 @@ async def send_json(self, data: typing.Any, mode: str = "text") -> None:
         else:
             await self.send({"type": "websocket.send", "bytes": text.encode("utf-8")})
 
-    async def close(self, code: int = 1000) -> None:
-        await self.send({"type": "websocket.close", "code": code})
+    async def close(self, code: int = 1000, reason: str = None) -> None:
+        await self.send(
+            {"type": "websocket.close", "code": code, "reason": reason or ""}
+        )
 
 
 class WebSocketClose:
-    def __init__(self, code: int = 1000) -> None:
+    def __init__(self, code: int = 1000, reason: str = None) -> None:
         self.code = code
+        self.reason = reason or ""
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        await send({"type": "websocket.close", "code": self.code})
+        await send(
+            {"type": "websocket.close", "code": self.code, "reason": self.reason}
+        )


Test patch:
diff --git a/tests/test_websockets.py b/tests/test_websockets.py
--- a/tests/test_websockets.py
+++ b/tests/test_websockets.py
@@ -405,3 +405,20 @@ async def mock_send(message):
     assert websocket == websocket
     assert websocket in {websocket}
     assert {websocket} == {websocket}
+
+
+def test_websocket_close_reason(test_client_factory) -> None:
+    def app(scope):
+        async def asgi(receive, send):
+            websocket = WebSocket(scope, receive=receive, send=send)
+            await websocket.accept()
+            await websocket.close(code=status.WS_1001_GOING_AWAY, reason="Going Away")
+
+        return asgi
+
+    client = test_client_factory(app)
+    with client.websocket_connect("/") as websocket:
+        with pytest.raises(WebSocketDisconnect) as exc:
+            websocket.receive_text()
+        assert exc.value.code == status.WS_1001_GOING_AWAY
+        assert exc.value.reason == "Going Away"