James Kaplan
6 years ago
committed by
Sebastián Ramírez
Sebastián Ramírez
12 changed files with 305 additions and 17 deletions
@ -0,0 +1,78 @@ |
|||||
|
from fastapi import Cookie, Depends, FastAPI, Header |
||||
|
from starlette.responses import HTMLResponse |
||||
|
from starlette.status import WS_1008_POLICY_VIOLATION |
||||
|
from starlette.websockets import WebSocket |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
html = """ |
||||
|
<!DOCTYPE html> |
||||
|
<html> |
||||
|
<head> |
||||
|
<title>Chat</title> |
||||
|
</head> |
||||
|
<body> |
||||
|
<h1>WebSocket Chat</h1> |
||||
|
<form action="" onsubmit="sendMessage(event)"> |
||||
|
<label>Item ID: <input type="text" id="itemId" autocomplete="off" value="foo"/></label> |
||||
|
<button onclick="connect(event)">Connect</button> |
||||
|
<br> |
||||
|
<label>Message: <input type="text" id="messageText" autocomplete="off"/></label> |
||||
|
<button>Send</button> |
||||
|
</form> |
||||
|
<ul id='messages'> |
||||
|
</ul> |
||||
|
<script> |
||||
|
var ws = null; |
||||
|
function connect(event) { |
||||
|
var input = document.getElementById("itemId") |
||||
|
ws = new WebSocket("ws://localhost:8000/items/" + input.value + "/ws"); |
||||
|
ws.onmessage = function(event) { |
||||
|
var messages = document.getElementById('messages') |
||||
|
var message = document.createElement('li') |
||||
|
var content = document.createTextNode(event.data) |
||||
|
message.appendChild(content) |
||||
|
messages.appendChild(message) |
||||
|
}; |
||||
|
} |
||||
|
function sendMessage(event) { |
||||
|
var input = document.getElementById("messageText") |
||||
|
ws.send(input.value) |
||||
|
input.value = '' |
||||
|
event.preventDefault() |
||||
|
} |
||||
|
</script> |
||||
|
</body> |
||||
|
</html> |
||||
|
""" |
||||
|
|
||||
|
|
||||
|
@app.get("/") |
||||
|
async def get(): |
||||
|
return HTMLResponse(html) |
||||
|
|
||||
|
|
||||
|
async def get_cookie_or_client( |
||||
|
websocket: WebSocket, session: str = Cookie(None), x_client: str = Header(None) |
||||
|
): |
||||
|
if session is None and x_client is None: |
||||
|
await websocket.close(code=WS_1008_POLICY_VIOLATION) |
||||
|
return session or x_client |
||||
|
|
||||
|
|
||||
|
@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), |
||||
|
): |
||||
|
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}" |
||||
|
) |
||||
|
if q is not None: |
||||
|
await websocket.send_text(f"Query parameter q is: {q}") |
||||
|
await websocket.send_text(f"Message text was: {data}, for item ID: {item_id}") |
||||
@ -0,0 +1,25 @@ |
|||||
|
import pytest |
||||
|
from starlette.testclient import TestClient |
||||
|
from starlette.websockets import WebSocketDisconnect |
||||
|
from websockets.tutorial001 import app |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
|
||||
|
|
||||
|
def test_main(): |
||||
|
response = client.get("/") |
||||
|
assert response.status_code == 200 |
||||
|
assert b"<!DOCTYPE html>" in response.content |
||||
|
|
||||
|
|
||||
|
def test_websocket(): |
||||
|
with pytest.raises(WebSocketDisconnect): |
||||
|
with client.websocket_connect("/ws") as websocket: |
||||
|
message = "Message one" |
||||
|
websocket.send_text(message) |
||||
|
data = websocket.receive_text() |
||||
|
assert data == f"Message text was: {message}" |
||||
|
message = "Message two" |
||||
|
websocket.send_text(message) |
||||
|
data = websocket.receive_text() |
||||
|
assert data == f"Message text was: {message}" |
||||
@ -0,0 +1,83 @@ |
|||||
|
import pytest |
||||
|
from starlette.testclient import TestClient |
||||
|
from starlette.websockets import WebSocketDisconnect |
||||
|
from websockets.tutorial002 import app |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
|
||||
|
|
||||
|
def test_main(): |
||||
|
response = client.get("/") |
||||
|
assert response.status_code == 200 |
||||
|
assert b"<!DOCTYPE html>" in response.content |
||||
|
|
||||
|
|
||||
|
def test_websocket_with_cookie(): |
||||
|
with pytest.raises(WebSocketDisconnect): |
||||
|
with client.websocket_connect( |
||||
|
"/items/1/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" |
||||
|
data = websocket.receive_text() |
||||
|
assert data == f"Message text was: {message}, for item ID: 1" |
||||
|
message = "Message two" |
||||
|
websocket.send_text(message) |
||||
|
data = websocket.receive_text() |
||||
|
assert data == "Session Cookie or X-Client Header value is: fakesession" |
||||
|
data = websocket.receive_text() |
||||
|
assert data == f"Message text was: {message}, for item ID: 1" |
||||
|
|
||||
|
|
||||
|
def test_websocket_with_header(): |
||||
|
with pytest.raises(WebSocketDisconnect): |
||||
|
with client.websocket_connect( |
||||
|
"/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" |
||||
|
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" |
||||
|
data = websocket.receive_text() |
||||
|
assert data == f"Message text was: {message}, for item ID: 2" |
||||
|
|
||||
|
|
||||
|
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: |
||||
|
message = "Message one" |
||||
|
websocket.send_text(message) |
||||
|
data = websocket.receive_text() |
||||
|
assert data == "Session Cookie or X-Client Header value is: xmen" |
||||
|
data = websocket.receive_text() |
||||
|
assert data == "Query parameter q is: baz" |
||||
|
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" |
||||
|
data = websocket.receive_text() |
||||
|
assert data == "Query parameter q is: baz" |
||||
|
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") |
||||
|
|
||||
|
|
||||
|
def test_websocket_invalid_data(): |
||||
|
with pytest.raises(WebSocketDisconnect): |
||||
|
client.websocket_connect("/items/foo/ws", headers={"X-Client": "xmen"}) |
||||
Loading…
Reference in new issue