Repository: django/channels
Pull request #2002 (2023-04-07 13:54:25)

Updated websocket consumers for ASGI v2.3. 
This PR adds the "headers" field for `websocket.accept` messages and the "reason" field for `websocket.close`, which
were added in ASGI spec versions 2.1 and 2.3 respectively.

kristjanvalur: I've updated, could you re-run the workflow?
kristjanvalur: Oops, sorry about that, everything should now pass, including the linting.
kristjanvalur: I did the fixes with github.dev, could not test them, but I´m crossing my fingers.

Patch:
diff --git a/channels/generic/websocket.py b/channels/generic/websocket.py
index 6c8ca57..9ce2657 100644
--- a/channels/generic/websocket.py
+++ b/channels/generic/websocket.py
@@ -44,11 +44,15 @@ class WebsocketConsumer(SyncConsumer):
     def connect(self):
         self.accept()
 
-    def accept(self, subprotocol=None):
+    def accept(self, subprotocol=None, headers=None):
         """
         Accepts an incoming socket
         """
-        super().send({"type": "websocket.accept", "subprotocol": subprotocol})
+        message = {"type": "websocket.accept", "subprotocol": subprotocol}
+        if headers:
+            message["headers"] = list(headers)
+
+        super().send(message)
 
     def websocket_receive(self, message):
         """
@@ -79,14 +83,16 @@ class WebsocketConsumer(SyncConsumer):
         if close:
             self.close(close)
 
-    def close(self, code=None):
+    def close(self, code=None, reason=None):
         """
         Closes the WebSocket from the server end
         """
+        message = {"type": "websocket.close"}
         if code is not None and code is not True:
-            super().send({"type": "websocket.close", "code": code})
-        else:
-            super().send({"type": "websocket.close"})
+            message["code"] = code
+        if reason:
+            message["reason"] = reason
+        super().send(message)
 
     def websocket_disconnect(self, message):
         """
@@ -179,11 +185,14 @@ class AsyncWebsocketConsumer(AsyncConsumer):
     async def connect(self):
         await self.accept()
 
-    async def accept(self, subprotocol=None):
+    async def accept(self, subprotocol=None, headers=None):
         """
         Accepts an incoming socket
         """
-        await super().send({"type": "websocket.accept", "subprotocol": subprotocol})
+        message = {"type": "websocket.accept", "subprotocol": subprotocol}
+        if headers:
+            message["headers"] = list(headers)
+        await super().send(message)
 
     async def websocket_receive(self, message):
         """
@@ -214,14 +223,16 @@ class AsyncWebsocketConsumer(AsyncConsumer):
         if close:
             await self.close(close)
 
-    async def close(self, code=None):
+    async def close(self, code=None, reason=None):
         """
         Closes the WebSocket from the server end
         """
+        message = {"type": "websocket.close"}
         if code is not None and code is not True:
-            await super().send({"type": "websocket.close", "code": code})
-        else:
-            await super().send({"type": "websocket.close"})
+            message["code"] = code
+        if reason:
+            message["reason"] = reason
+        await super().send(message)
 
     async def websocket_disconnect(self, message):
         """


Test patch:
diff --git a/channels/testing/websocket.py b/channels/testing/websocket.py
index e176321..dd48686 100644
--- a/channels/testing/websocket.py
+++ b/channels/testing/websocket.py
@@ -12,7 +12,9 @@ class WebsocketCommunicator(ApplicationCommunicator):
     (uninstantiated) along with the initial connection parameters.
     """
 
-    def __init__(self, application, path, headers=None, subprotocols=None):
+    def __init__(
+        self, application, path, headers=None, subprotocols=None, spec_version=None
+    ):
         if not isinstance(path, str):
             raise TypeError("Expected str, got {}".format(type(path)))
         parsed = urlparse(path)
@@ -23,7 +25,10 @@ class WebsocketCommunicator(ApplicationCommunicator):
             "headers": headers or [],
             "subprotocols": subprotocols or [],
         }
+        if spec_version:
+            self.scope["spec_version"] = spec_version
         super().__init__(application, self.scope)
+        self.response_headers = None
 
     async def connect(self, timeout=1):
         """
@@ -37,6 +42,8 @@ class WebsocketCommunicator(ApplicationCommunicator):
         if response["type"] == "websocket.close":
             return (False, response.get("code", 1000))
         else:
+            assert response["type"] == "websocket.accept"
+            self.response_headers = response.get("headers", [])
             return (True, response.get("subprotocol", None))
 
     async def send_to(self, text_data=None, bytes_data=None):
diff --git a/tests/test_generic_websocket.py b/tests/test_generic_websocket.py
index 10a761a..73cdb48 100644
--- a/tests/test_generic_websocket.py
+++ b/tests/test_generic_websocket.py
@@ -424,3 +424,57 @@ async def test_block_leading_dot_type_function_call():
             ValueError, match=r"Malformed type in message \(leading underscore\)"
         ):
             await communicator.receive_from()
+
+
+@pytest.mark.parametrize("async_consumer", [False, True])
+@pytest.mark.django_db
+@pytest.mark.asyncio
+async def test_accept_headers(async_consumer):
+    """
+    Tests that JsonWebsocketConsumer is implemented correctly.
+    """
+
+    class TestConsumer(WebsocketConsumer):
+        def connect(self):
+            self.accept(headers=[[b"foo", b"bar"]])
+
+    class AsyncTestConsumer(AsyncWebsocketConsumer):
+        async def connect(self):
+            await self.accept(headers=[[b"foo", b"bar"]])
+
+    app = AsyncTestConsumer() if async_consumer else TestConsumer()
+
+    # Open a connection
+    communicator = WebsocketCommunicator(app, "/testws/", spec_version="2.3")
+    connected, _ = await communicator.connect()
+    assert connected
+    assert communicator.response_headers == [[b"foo", b"bar"]]
+
+
+@pytest.mark.parametrize("async_consumer", [False, True])
+@pytest.mark.django_db
+@pytest.mark.asyncio
+async def test_close_reason(async_consumer):
+    """
+    Tests that JsonWebsocketConsumer is implemented correctly.
+    """
+
+    class TestConsumer(WebsocketConsumer):
+        def connect(self):
+            self.accept()
+            self.close(code=4007, reason="test reason")
+
+    class AsyncTestConsumer(AsyncWebsocketConsumer):
+        async def connect(self):
+            await self.accept()
+            await self.close(code=4007, reason="test reason")
+
+    app = AsyncTestConsumer() if async_consumer else TestConsumer()
+
+    # Open a connection
+    communicator = WebsocketCommunicator(app, "/testws/", spec_version="2.3")
+    connected, _ = await communicator.connect()
+    msg = await communicator.receive_output()
+    assert msg["type"] == "websocket.close"
+    assert msg["code"] == 4007
+    assert msg["reason"] == "test reason"