forked from encode/starlette
/
test_testclient.py
226 lines (168 loc) Β· 6.78 KB
/
test_testclient.py
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
import asyncio
import itertools
import sys
import anyio
import pytest
import sniffio
import trio.lowlevel
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.responses import JSONResponse
from starlette.websockets import WebSocket, WebSocketDisconnect
if sys.version_info >= (3, 7):
from asyncio import current_task as asyncio_current_task # pragma: no cover
else:
asyncio_current_task = asyncio.Task.current_task # pragma: no cover
mock_service = Starlette()
@mock_service.route("/")
def mock_service_endpoint(request):
return JSONResponse({"mock": "example"})
_identity_runvar: anyio.lowlevel.RunVar[int] = anyio.lowlevel.RunVar("_identity_runvar")
def get_identity(counter):
try:
return _identity_runvar.get()
except LookupError:
token = next(counter)
_identity_runvar.set(token)
return token
def current_task():
# anyio's TaskInfo comparisons are invalid after their associated native
# task object is GC'd https://github.com/agronholm/anyio/issues/324
asynclib_name = sniffio.current_async_library()
if asynclib_name == "trio":
return trio.lowlevel.current_task()
if asynclib_name == "asyncio":
task = asyncio_current_task()
if task is None:
raise RuntimeError("must be called from a running task") # pragma: no cover
return task
raise RuntimeError(f"unsupported asynclib={asynclib_name}") # pragma: no cover
def create_app(test_client_factory, counter=itertools.count()):
app = Starlette()
@app.on_event("startup")
async def get_startup_thread():
app.startup_task = current_task()
app.startup_loop = get_identity(counter)
@app.on_event("shutdown")
async def get_shutdown_thread():
app.shutdown_task = current_task()
app.shutdown_loop = get_identity(counter)
@app.route("/")
def homepage(request):
client = test_client_factory(mock_service)
response = client.get("/")
return JSONResponse(response.json())
@app.route("/thread")
async def thread(request):
return JSONResponse(get_identity(counter))
return app
startup_error_app = Starlette()
@startup_error_app.on_event("startup")
def startup():
raise RuntimeError()
def test_use_testclient_in_endpoint(test_client_factory):
"""
We should be able to use the test client within applications.
This is useful if we need to mock out other services,
during tests or in development.
"""
client = test_client_factory(create_app(test_client_factory))
response = client.get("/")
assert response.json() == {"mock": "example"}
def test_use_testclient_as_contextmanager(test_client_factory, anyio_backend_name):
"""
This test asserts a number of properties that are important for an
app level task_group
"""
app = create_app(test_client_factory, counter=itertools.count())
client = test_client_factory(app)
with client:
# within a TestClient context every async request runs in the same thread
assert client.get("/thread").json() == 0
assert client.get("/thread").json() == 0
# that thread is also the same as the lifespan thread
assert app.startup_loop == 0
assert app.shutdown_loop == 0
# lifespan events run in the same task, this is important because a task
# group must be entered and exited in the same task.
assert app.startup_task is app.shutdown_task
# outside the TestClient context, new requests continue to spawn in new
# eventloops in new threads
assert client.get("/thread").json() == 1
assert client.get("/thread").json() == 2
first_task = app.startup_task
with client:
# the TestClient context can be re-used, starting a new lifespan task
# in a new thread
assert client.get("/thread").json() == 3
assert client.get("/thread").json() == 3
assert app.startup_loop == 3
assert app.shutdown_loop == 3
# lifespan events still run in the same task, with the context but...
assert app.startup_task is app.shutdown_task
# ... the second TestClient context creates a new lifespan task.
assert first_task is not app.startup_task
def test_error_on_startup(test_client_factory):
with pytest.raises(RuntimeError):
with test_client_factory(startup_error_app):
pass # pragma: no cover
def test_exception_in_middleware(test_client_factory):
class MiddlewareException(Exception):
pass
class BrokenMiddleware:
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
raise MiddlewareException()
broken_middleware = Starlette(middleware=[Middleware(BrokenMiddleware)])
with pytest.raises(MiddlewareException):
with test_client_factory(broken_middleware):
pass # pragma: no cover
def test_testclient_asgi2(test_client_factory):
def app(scope):
async def inner(receive, send):
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"text/plain"]],
}
)
await send({"type": "http.response.body", "body": b"Hello, world!"})
return inner
client = test_client_factory(app)
response = client.get("/")
assert response.text == "Hello, world!"
def test_testclient_asgi3(test_client_factory):
async def app(scope, receive, send):
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"text/plain"]],
}
)
await send({"type": "http.response.body", "body": b"Hello, world!"})
client = test_client_factory(app)
response = client.get("/")
assert response.text == "Hello, world!"
def test_websocket_blocking_receive(test_client_factory):
def app(scope):
async def respond(websocket):
await websocket.send_json({"message": "test"})
async def asgi(receive, send):
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
async with anyio.create_task_group() as task_group:
task_group.start_soon(respond, websocket)
try:
# this will block as the client does not send us data
# it should not prevent `respond` from executing though
await websocket.receive_json()
except WebSocketDisconnect:
pass
return asgi
client = test_client_factory(app)
with client.websocket_connect("/") as websocket:
data = websocket.receive_json()
assert data == {"message": "test"}