Additional headers for WS accept message.

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
 

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:

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":

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):