Additional headers for WS accept message.

matiuszka · matiuszka · commit a080bf5d05d8 · 2021-12-14T11:17:35.000+01:00
diff --git a/docs/websockets.md b/docs/websockets.md
@@ -1,4 +1,3 @@
-
 Starlette includes a `WebSocket` class that fulfils a similar role
 to the HTTP request, but that allows sending and receiving data on a websocket.
 
@@ -51,7 +50,7 @@ For example: `websocket.path_params['username']`
 
 ### Accepting the connection
 
-* `await websocket.accept(subprotocol=None)`
+* `await websocket.accept(subprotocol=None, headers=None)`
 
 ### Sending data
 
diff --git a/starlette/testclient.py b/starlette/testclient.py
@@ -296,6 +296,7 @@ def __init__(
         self.app = app
         self.scope = scope
         self.accepted_subprotocol = None
+        self.additional_headers = None
         self.portal_factory = portal_factory
         self._receive_queue: "queue.Queue[typing.Any]" = queue.Queue()
         self._send_queue: "queue.Queue[typing.Any]" = queue.Queue()
@@ -313,6 +314,7 @@ def __enter__(self) -> "WebSocketTestSession":
             self.exit_stack.close()
             raise
         self.accepted_subprotocol = message.get("subprotocol", None)
+        self.additional_headers = message.get("headers", None)
         return self
 
     def __exit__(self, *args: typing.Any) -> None:
diff --git a/starlette/websockets.py b/starlette/websockets.py
@@ -69,11 +69,17 @@ async def send(self, message: Message) -> None:
         else:
             raise RuntimeError('Cannot call "send" once a close message has been sent.')
 
-    async def accept(self, subprotocol: str = None) -> None:
+    async def accept(
+        self,
+        subprotocol: str = None,
+        headers: typing.Iterable[typing.Tuple[bytes, bytes]] = None,
+    ) -> None:
         if self.client_state == WebSocketState.CONNECTING:
             # If we haven't yet seen the 'connect' message, then wait for it first.
             await self.receive()
-        await self.send({"type": "websocket.accept", "subprotocol": subprotocol})
+        await self.send(
+            {"type": "websocket.accept", "subprotocol": subprotocol, "headers": headers}
+        )
 
     def _raise_on_disconnect(self, message: Message) -> None:
         if message["type"] == "websocket.disconnect":
diff --git a/tests/test_websockets.py b/tests/test_websockets.py
@@ -301,6 +301,20 @@ async def asgi(receive, send):
         assert websocket.accepted_subprotocol == "wamp"
 
 
+def test_additional_headers(test_client_factory):
+    def app(scope):
+        async def asgi(receive, send):
+            websocket = WebSocket(scope, receive=receive, send=send)
+            await websocket.accept(headers=[(b"additional", b"header")])
+            await websocket.close()
+
+        return asgi
+
+    client = test_client_factory(app)
+    with client.websocket_connect("/") as websocket:
+        websocket.additional_headers = [(b"additional", b"header")]
+
+
 def test_websocket_exception(test_client_factory):
     def app(scope):
         async def asgi(receive, send):
