📝 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]>
5 years ago · e4300769ac
4 changed files with 79 additions and 52 deletions
--- a/docs/en/docs/advanced/websockets.md
+++ b/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
-
-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
+### Try the WebSockets with dependencies

 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">
-
-You can send (and receive) many messages:
+## More info

-<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>.
--- a/docs/en/docs/img/tutorial/websockets/image05.png
+++ b/docs/en/docs/img/tutorial/websockets/image05.png
--- a/docs_src/websockets/tutorial002.py
+++ b/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")
-                ws = new WebSocket("ws://localhost:8000/items/" + input.value + "/ws");
+                var itemId = document.getElementById("itemId")
+                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(
-    websocket: WebSocket, session: str = Cookie(None), x_client: str = Header(None)
+async def get_cookie_or_token(
+    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,
-    q: str = None,
-    cookie_or_client: str = Depends(get_cookie_or_client),
+    item_id: str,
+    q: int = None,
+    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}")
--- a/tests/test_tutorial/test_websockets/test_tutorial002.py
+++ b/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(
-            "/items/2/ws", headers={"X-Client": "xmen"}
-        ) as websocket:
+        with client.websocket_connect("/items/bar/ws?token=some-token") 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(
-            "/items/2/ws?q=baz", headers={"X-Client": "xmen"}
-        ) as websocket:
+        with client.websocket_connect("/items/2/ws?q=3&token=some-token") 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")