[Feature Request] Inheriting cookies from TestClient to WebSocketConnection.

[otsuka]

First of all, thanks for developing such a useful library.

I use cookies for authentication, and accept WebSocket connection from the authenticated user. I can set cookies to a TestClient instance as bellow, although, the cookies are not included in WebSocket connection request.

TestClient.websocket_connect() method has extra_headers argument, but adding cookies directly to HTTP Header is not easy.

So It would be helpful if WebSocketSession could take over cookies from TestClient.

async with TestClient(app) as client
    ck = SimpleCookie()
    ck["foo"] = "bar"
    client.cookie_jar = ck

    async with client.websocket_connect(endpoint) as ws:
        ...

[masipcat]

Hi @otsuka, Feel free to open a PR with these features if you want :) but if you prefer, I can do it myself during the weekend.

[otsuka]

@masipcat I will appreciate it if you would do it.

[otsuka]

In my use case, I have to write tests dealing with multiple users like chat app.

So it might be convenient to be able to set cookies in each WebSocketSession instance in addition to inheriting cookies from TestClient.

It would be useful if we could set cookies as easily as the code below (like aiohttp.ClientSession)

async with TestClient(app):  # invoking 'startup' only once

    # I don't use `async with WebSocketSession` to avoid deep nesting when handling multiple ws connections
    ws1 = WebSocketSession(app, ws_endpoint, cookies={'sessionid': 'xxxxxxxxxxxxxx'})
    ws2 = WebSocketSession(app, ws_endpoint, cookies={'sessionid': 'yyyyyyyyyyyyyy'})
    await ws1.connect()
    await ws2.connect()

    await ws1.send_str('Hello')  # BTW, send_text() is better for consistency, isn't it?
    text = await ws2.receive_text()
    assert text == 'Hello'

    await ws2.close()
    await ws1.close()

[masipcat]

Version 1.4.0 released! I added the argument 'cookies' to websocket_connect() (and WebSocketSession as well) and I renamed ws.send_str() to ws.send_text().