1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
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
|
