mirror of https://github.com/tiangolo/fastapi.git
108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
import importlib
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from fastapi.websockets import WebSocketDisconnect
|
|
|
|
from ...utils import needs_py39, needs_py310
|
|
|
|
|
|
@pytest.fixture(
|
|
name="app",
|
|
params=[
|
|
"tutorial002",
|
|
pytest.param("tutorial002_py310", marks=needs_py310),
|
|
"tutorial002_an",
|
|
pytest.param("tutorial002_an_py39", marks=needs_py39),
|
|
pytest.param("tutorial002_an_py310", marks=needs_py310),
|
|
],
|
|
)
|
|
def get_app(request: pytest.FixtureRequest):
|
|
mod = importlib.import_module(f"docs_src.websockets.{request.param}")
|
|
|
|
return mod.app
|
|
|
|
|
|
def test_main(app: FastAPI):
|
|
client = TestClient(app)
|
|
response = client.get("/")
|
|
assert response.status_code == 200, response.text
|
|
assert b"<!DOCTYPE html>" in response.content
|
|
|
|
|
|
def test_websocket_with_cookie(app: FastAPI):
|
|
client = TestClient(app, cookies={"session": "fakesession"})
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with client.websocket_connect("/items/foo/ws") as websocket:
|
|
message = "Message one"
|
|
websocket.send_text(message)
|
|
data = websocket.receive_text()
|
|
assert data == "Session cookie or query token value is: fakesession"
|
|
data = websocket.receive_text()
|
|
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 query token value is: fakesession"
|
|
data = websocket.receive_text()
|
|
assert data == f"Message text was: {message}, for item ID: foo"
|
|
|
|
|
|
def test_websocket_with_header(app: FastAPI):
|
|
client = TestClient(app)
|
|
with pytest.raises(WebSocketDisconnect):
|
|
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 query token value is: some-token"
|
|
data = websocket.receive_text()
|
|
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 query token value is: some-token"
|
|
data = websocket.receive_text()
|
|
assert data == f"Message text was: {message}, for item ID: bar"
|
|
|
|
|
|
def test_websocket_with_header_and_query(app: FastAPI):
|
|
client = TestClient(app)
|
|
with pytest.raises(WebSocketDisconnect):
|
|
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 query token value is: some-token"
|
|
data = websocket.receive_text()
|
|
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 query token value is: some-token"
|
|
data = websocket.receive_text()
|
|
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(app: FastAPI):
|
|
client = TestClient(app)
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with client.websocket_connect("/items/foo/ws"):
|
|
pytest.fail(
|
|
"did not raise WebSocketDisconnect on __enter__"
|
|
) # pragma: no cover
|
|
|
|
|
|
def test_websocket_invalid_data(app: FastAPI):
|
|
client = TestClient(app)
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with client.websocket_connect("/items/foo/ws?q=bar&token=some-token"):
|
|
pytest.fail(
|
|
"did not raise WebSocketDisconnect on __enter__"
|
|
) # pragma: no cover
|