happyz
/
fastapi
mirror of https://github.com/tiangolo/fastapi.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

📝 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 <Sean@Sean-Mac.local>
Co-authored-by: Sebastián Ramírez <tiangolo@gmail.com>
This commit is contained in:
Chih Sean Hsu 2020-06-14 04:51:34 +08:00 committed by GitHub
parent c6dd627bdd
commit e4300769ac
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 92 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

92
docs/en/docs/advanced/websockets.md
View File

@ -51,38 +51,7 @@ In your WebSocket route you can `await` for messages and send messages.
You can receive and send binary, text, and JSON data.
## Using `Depends` and others
In WebSocket endpoints you can import from `fastapi` and use:
* `Depends`
* `Security`
* `Cookie`
* `Header`
* `Path`
* `Query`
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"
{!../../../docs_src/websockets/tutorial002.py!}
```
!!! info
In a WebSocket it doesn't really make sense to raise an `HTTPException`. So it's better to close the WebSocket connection directly.
You can use a closing code from the <a href="https://tools.ietf.org/html/rfc6455#section-7.4.1" class="external-link" target="_blank">valid codes defined in the specification</a>.
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 it
If your file is named `main.py`, run your application with:
@ -115,3 +84,62 @@ 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:
* `Depends`
* `Security`
* `Cookie`
* `Header`
* `Path`
* `Query`
They work the same way as for other FastAPI endpoints/*path operations*:
```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!}
```
!!! info
In a WebSocket it doesn't really make sense to raise an `HTTPException`. So it's better to close the WebSocket connection directly.
You can use a closing code from the <a href="https://tools.ietf.org/html/rfc6455#section-7.4.1" class="external-link" target="_blank">valid codes defined in the specification</a>.
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.
### Try the WebSockets with dependencies
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>.
There you can set:
* The "Item ID", used in the path.
* The "Token" used as a query parameter.
!!! tip
Notice that the query `token` will be handled by a dependency.
With that you can connect the WebSocket and then send and receive messages:
<img src="/img/tutorial/websockets/image05.png">
## 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>.

BIN
docs/en/docs/img/tutorial/websockets/image05.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 52 KiB

27
docs_src/websockets/tutorial002.py
View File

@ -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}")

38
tests/test_tutorial/test_websockets/test_tutorial002.py
View File

@ -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")