📝 Update tutorial for WebSockets with dependencies (#1540)

* fix websockets/tutorial002.py

* fix tutorial002 in ws to correspond with test case

* reformat websocket tutorial002

* fix websocket tutorial002 coverage

* 📝 Update example for WebSockets with Depends

* ✅ Update and refactor tests for WebSockets with dependencies

* 👷 Trigger Travis, as it's not reporting to Codecov

* ✅ Update WebSocket tests to raise coverage

Co-authored-by: Chih Sean Hsu <[email protected]>
Co-authored-by: Sebastián Ramírez <[email protected]>

docs/en/docs/advanced/websockets.md

@@ -51,6 +51,40 @@ In your WebSocket route you can `await` for messages and send messages.
 
 You can receive and send binary, text, and JSON data.
 
+## Try it
+
+If your file is named `main.py`, run your application with:
+
+<div class="termy">
+
+```console
+$ uvicorn main:app --reload
+
+<span style="color: green;">INFO</span>:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
+```
+
+</div>
+
+Open your browser at <a href="http://127.0.0.1:8000" class="external-link" target="_blank">http://127.0.0.1:8000</a>.
+
+You will see a simple page like:
+
+<img src="/img/tutorial/websockets/image01.png">
+
+You can type messages in the input box, and send them:
+
+<img src="/img/tutorial/websockets/image02.png">
+
+And your **FastAPI** application with WebSockets will respond back:
+
+<img src="/img/tutorial/websockets/image03.png">
+
+You can send (and receive) many messages:
+
+<img src="/img/tutorial/websockets/image04.png">
+
+And all of them will use the same WebSocket connection.
+
 ## Using `Depends` and others
 
 In WebSocket endpoints you can import from `fastapi` and use:
@@ -64,7 +98,7 @@ In WebSocket endpoints you can import from `fastapi` and use:
 
 They work the same way as for other FastAPI endpoints/*path operations*:
 
-```Python hl_lines="53 54 55 56 57 58  61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76"
+```Python hl_lines="56 57 58 59 60 61  64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79"
 {!../../../docs_src/websockets/tutorial002.py!}
 ```
 
@@ -75,14 +109,7 @@ They work the same way as for other FastAPI endpoints/*path operations*:
 
     In the future, there will be a `WebSocketException` that you will be able to `raise` from anywhere, and add exception handlers for it. It depends on the <a href="https://github.com/encode/starlette/pull/527" class="external-link" target="_blank">PR #527</a> in Starlette.
 
-## More info
+### Try the WebSockets with dependencies
-
-To learn more about the options, check Starlette's documentation for:
-
-* <a href="https://www.starlette.io/websockets/" class="external-link" target="_blank">The `WebSocket` class</a>.
-* <a href="https://www.starlette.io/endpoints/#websocketendpoint" class="external-link" target="_blank">Class-based WebSocket handling</a>.
-
-## Test it
 
 If your file is named `main.py`, run your application with:
 
@@ -98,20 +125,21 @@ $ uvicorn main:app --reload
 
 Open your browser at <a href="http://127.0.0.1:8000" class="external-link" target="_blank">http://127.0.0.1:8000</a>.
 
-You will see a simple page like:
+There you can set:
 
-<img src="/img/tutorial/websockets/image01.png">
+* The "Item ID", used in the path.
+* The "Token" used as a query parameter.
 
-You can type messages in the input box, and send them:
+!!! tip
+    Notice that the query `token` will be handled by a dependency.
 
-<img src="/img/tutorial/websockets/image02.png">
+With that you can connect the WebSocket and then send and receive messages:
 
-And your **FastAPI** application with WebSockets will respond back:
+<img src="/img/tutorial/websockets/image05.png">
 
-<img src="/img/tutorial/websockets/image03.png">
+## More info
-
-You can send (and receive) many messages:
 
-<img src="/img/tutorial/websockets/image04.png">
+To learn more about the options, check Starlette's documentation for:
 
-And all of them will use the same WebSocket connection.
+* <a href="https://www.starlette.io/websockets/" class="external-link" target="_blank">The `WebSocket` class</a>.
+* <a href="https://www.starlette.io/endpoints/#websocketendpoint" class="external-link" target="_blank">Class-based WebSocket handling</a>.

docs_src/websockets/tutorial002.py

@@ -1,4 +1,4 @@
-from fastapi import Cookie, Depends, FastAPI, Header, WebSocket, status
+from fastapi import Cookie, Depends, FastAPI, Query, WebSocket, status
 from fastapi.responses import HTMLResponse
 
 app = FastAPI()
@@ -13,8 +13,9 @@ html = """
         <h1>WebSocket Chat</h1>
         <form action="" onsubmit="sendMessage(event)">
             <label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label>
+            <label>Token: <input type="text" id="token" autocomplete="off" value="some-key-token"/></label>
             <button onclick="connect(event)">Connect</button>
-            <br>
+            <hr>
             <label>Message: <input type="text" id="messageText" autocomplete="off"/></label>
             <button>Send</button>
         </form>
@@ -23,8 +24,9 @@ html = """
         <script>
         var ws = null;
             function connect(event) {
-                var input = document.getElementById("itemId")
+                var itemId = document.getElementById("itemId")
-                ws = new WebSocket("ws://localhost:8000/items/" + input.value + "/ws");
+                var token = document.getElementById("token")
+                ws = new WebSocket("ws://localhost:8000/items/" + itemId.value + "/ws?token=" + token.value);
                 ws.onmessage = function(event) {
                     var messages = document.getElementById('messages')
                     var message = document.createElement('li')
@@ -32,6 +34,7 @@ html = """
                     message.appendChild(content)
                     messages.appendChild(message)
                 };
+                event.preventDefault()
             }
             function sendMessage(event) {
                 var input = document.getElementById("messageText")
@@ -50,26 +53,26 @@ async def get():
     return HTMLResponse(html)
 
 
-async def get_cookie_or_client(
+async def get_cookie_or_token(
-    websocket: WebSocket, session: str = Cookie(None), x_client: str = Header(None)
+    websocket: WebSocket, session: str = Cookie(None), token: str = Query(None)
 ):
-    if session is None and x_client is None:
+    if session is None and token is None:
         await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
-    return session or x_client
+    return session or token
 
 
 @app.websocket("/items/{item_id}/ws")
 async def websocket_endpoint(
     websocket: WebSocket,
-    item_id: int,
+    item_id: str,
-    q: str = None,
+    q: int = None,
-    cookie_or_client: str = Depends(get_cookie_or_client),
+    cookie_or_token: str = Depends(get_cookie_or_token),
 ):
     await websocket.accept()
     while True:
         data = await websocket.receive_text()
         await websocket.send_text(
-            f"Session Cookie or X-Client Header value is: {cookie_or_client}"
+            f"Session cookie or query token value is: {cookie_or_token}"
         )
         if q is not None:
             await websocket.send_text(f"Query parameter q is: {q}")

tests/test_tutorial/test_websockets/test_tutorial002.py

@@ -15,69 +15,65 @@ def test_main():
 def test_websocket_with_cookie():
     with pytest.raises(WebSocketDisconnect):
         with client.websocket_connect(
-            "/items/1/ws", cookies={"session": "fakesession"}
+            "/items/foo/ws", cookies={"session": "fakesession"}
         ) as websocket:
             message = "Message one"
             websocket.send_text(message)
             data = websocket.receive_text()
-            assert data == "Session Cookie or X-Client Header value is: fakesession"
+            assert data == "Session cookie or query token value is: fakesession"
             data = websocket.receive_text()
-            assert data == f"Message text was: {message}, for item ID: 1"
+            assert data == f"Message text was: {message}, for item ID: foo"
             message = "Message two"
             websocket.send_text(message)
             data = websocket.receive_text()
-            assert data == "Session Cookie or X-Client Header value is: fakesession"
+            assert data == "Session cookie or query token value is: fakesession"
             data = websocket.receive_text()
-            assert data == f"Message text was: {message}, for item ID: 1"
+            assert data == f"Message text was: {message}, for item ID: foo"
 
 
 def test_websocket_with_header():
     with pytest.raises(WebSocketDisconnect):
-        with client.websocket_connect(
+        with client.websocket_connect("/items/bar/ws?token=some-token") as websocket:
-            "/items/2/ws", headers={"X-Client": "xmen"}
-        ) as websocket:
             message = "Message one"
             websocket.send_text(message)
             data = websocket.receive_text()
-            assert data == "Session Cookie or X-Client Header value is: xmen"
+            assert data == "Session cookie or query token value is: some-token"
             data = websocket.receive_text()
-            assert data == f"Message text was: {message}, for item ID: 2"
+            assert data == f"Message text was: {message}, for item ID: bar"
             message = "Message two"
             websocket.send_text(message)
             data = websocket.receive_text()
-            assert data == "Session Cookie or X-Client Header value is: xmen"
+            assert data == "Session cookie or query token value is: some-token"
             data = websocket.receive_text()
-            assert data == f"Message text was: {message}, for item ID: 2"
+            assert data == f"Message text was: {message}, for item ID: bar"
 
 
 def test_websocket_with_header_and_query():
     with pytest.raises(WebSocketDisconnect):
-        with client.websocket_connect(
+        with client.websocket_connect("/items/2/ws?q=3&token=some-token") as websocket:
-            "/items/2/ws?q=baz", headers={"X-Client": "xmen"}
-        ) as websocket:
             message = "Message one"
             websocket.send_text(message)
             data = websocket.receive_text()
-            assert data == "Session Cookie or X-Client Header value is: xmen"
+            assert data == "Session cookie or query token value is: some-token"
             data = websocket.receive_text()
-            assert data == "Query parameter q is: baz"
+            assert data == "Query parameter q is: 3"
             data = websocket.receive_text()
             assert data == f"Message text was: {message}, for item ID: 2"
             message = "Message two"
             websocket.send_text(message)
             data = websocket.receive_text()
-            assert data == "Session Cookie or X-Client Header value is: xmen"
+            assert data == "Session cookie or query token value is: some-token"
             data = websocket.receive_text()
-            assert data == "Query parameter q is: baz"
+            assert data == "Query parameter q is: 3"
             data = websocket.receive_text()
             assert data == f"Message text was: {message}, for item ID: 2"
 
 
 def test_websocket_no_credentials():
     with pytest.raises(WebSocketDisconnect):
-        client.websocket_connect("/items/2/ws")
+        client.websocket_connect("/items/foo/ws")
 
 
 def test_websocket_invalid_data():
     with pytest.raises(WebSocketDisconnect):
-        client.websocket_connect("/items/foo/ws", headers={"X-Client": "xmen"})
+        client.websocket_connect("/items/foo/ws?q=bar&token=some-token")