From 41b06d2378243ba7fee5ce28207d5ea316c2d6ac Mon Sep 17 00:00:00 2001
From: Marcelo Trylesinski <marcelotryle@gmail.com>
Date: Tue, 10 May 2022 07:44:10 +0200
Subject: [PATCH] Add reason to WebSocketException

---
 starlette/exceptions.py  | 7 ++++---
 tests/test_exceptions.py | 9 +++++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/starlette/exceptions.py b/starlette/exceptions.py
index 9eeab7933..6b173391a 100644
--- a/starlette/exceptions.py
+++ b/starlette/exceptions.py
@@ -28,12 +28,13 @@ def __repr__(self) -> str:
 
 
 class WebSocketException(Exception):
-    def __init__(self, code: int) -> None:
+    def __init__(self, code: int, reason: typing.Optional[str] = None) -> None:
         self.code = code
+        self.reason = reason or ""
 
     def __repr__(self) -> str:
         class_name = self.__class__.__name__
-        return f"{class_name}(code={self.code!r})"
+        return f"{class_name}(code={self.code!r}, reason={self.reason!r})"
 
 
 class ExceptionMiddleware:
@@ -133,4 +134,4 @@ def http_exception(self, request: Request, exc: HTTPException) -> Response:
     async def websocket_exception(
         self, websocket: WebSocket, exc: WebSocketException
     ) -> None:
-        await websocket.close(code=exc.code)
+        await websocket.close(code=exc.code, reason=exc.reason)
diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
index 7824f8bae..c61186e64 100644
--- a/tests/test_exceptions.py
+++ b/tests/test_exceptions.py
@@ -133,9 +133,14 @@ class CustomHTTPException(HTTPException):
 
 
 def test_websocket_repr():
-    assert repr(WebSocketException(1008)) == ("WebSocketException(code=1008)")
+    assert repr(WebSocketException(1008, reason="Policy Violation")) == (
+        "WebSocketException(code=1008, reason='Policy Violation')"
+    )
 
     class CustomWebSocketException(WebSocketException):
         pass
 
-    assert repr(CustomWebSocketException(1013)) == "CustomWebSocketException(code=1013)"
+    assert (
+        repr(CustomWebSocketException(1013, reason="Something custom"))
+        == "CustomWebSocketException(code=1013, reason='Something custom')"
+    )