fastapi

1
import json
2
from typing import Dict
3

4
import pytest
5
from fastapi import BackgroundTasks, Depends, FastAPI
6
from fastapi.responses import StreamingResponse
7
from fastapi.testclient import TestClient
8

9
app = FastAPI()
10
state = {
11
    "/async": "asyncgen not started",
12
    "/sync": "generator not started",
13
    "/async_raise": "asyncgen raise not started",
14
    "/sync_raise": "generator raise not started",
15
    "context_a": "not started a",
16
    "context_b": "not started b",
17
    "bg": "not set",
18
    "sync_bg": "not set",
19
}
20

21
errors = []
22

23

24
async def get_state():
25
    return state
26

27

28
class AsyncDependencyError(Exception):
29
    pass
30

31

32
class SyncDependencyError(Exception):
33
    pass
34

35

36
class OtherDependencyError(Exception):
37
    pass
38

39

40
async def asyncgen_state(state: Dict[str, str] = Depends(get_state)):
41
    state["/async"] = "asyncgen started"
42
    yield state["/async"]
43
    state["/async"] = "asyncgen completed"
44

45

46
def generator_state(state: Dict[str, str] = Depends(get_state)):
47
    state["/sync"] = "generator started"
48
    yield state["/sync"]
49
    state["/sync"] = "generator completed"
50

51

52
async def asyncgen_state_try(state: Dict[str, str] = Depends(get_state)):
53
    state["/async_raise"] = "asyncgen raise started"
54
    try:
55
        yield state["/async_raise"]
56
    except AsyncDependencyError:
57
        errors.append("/async_raise")
58
        raise
59
    finally:
60
        state["/async_raise"] = "asyncgen raise finalized"
61

62

63
def generator_state_try(state: Dict[str, str] = Depends(get_state)):
64
    state["/sync_raise"] = "generator raise started"
65
    try:
66
        yield state["/sync_raise"]
67
    except SyncDependencyError:
68
        errors.append("/sync_raise")
69
        raise
70
    finally:
71
        state["/sync_raise"] = "generator raise finalized"
72

73

74
async def context_a(state: dict = Depends(get_state)):
75
    state["context_a"] = "started a"
76
    try:
77
        yield state
78
    finally:
79
        state["context_a"] = "finished a"
80

81

82
async def context_b(state: dict = Depends(context_a)):
83
    state["context_b"] = "started b"
84
    try:
85
        yield state
86
    finally:
87
        state["context_b"] = f"finished b with a: {state['context_a']}"
88

89

90
@app.get("/async")
91
async def get_async(state: str = Depends(asyncgen_state)):
92
    return state
93

94

95
@app.get("/sync")
96
async def get_sync(state: str = Depends(generator_state)):
97
    return state
98

99

100
@app.get("/async_raise")
101
async def get_async_raise(state: str = Depends(asyncgen_state_try)):
102
    assert state == "asyncgen raise started"
103
    raise AsyncDependencyError()
104

105

106
@app.get("/sync_raise")
107
async def get_sync_raise(state: str = Depends(generator_state_try)):
108
    assert state == "generator raise started"
109
    raise SyncDependencyError()
110

111

112
@app.get("/async_raise_other")
113
async def get_async_raise_other(state: str = Depends(asyncgen_state_try)):
114
    assert state == "asyncgen raise started"
115
    raise OtherDependencyError()
116

117

118
@app.get("/sync_raise_other")
119
async def get_sync_raise_other(state: str = Depends(generator_state_try)):
120
    assert state == "generator raise started"
121
    raise OtherDependencyError()
122

123

124
@app.get("/context_b")
125
async def get_context_b(state: dict = Depends(context_b)):
126
    return state
127

128

129
@app.get("/context_b_raise")
130
async def get_context_b_raise(state: dict = Depends(context_b)):
131
    assert state["context_b"] == "started b"
132
    assert state["context_a"] == "started a"
133
    raise OtherDependencyError()
134

135

136
@app.get("/context_b_bg")
137
async def get_context_b_bg(tasks: BackgroundTasks, state: dict = Depends(context_b)):
138
    async def bg(state: dict):
139
        state["bg"] = f"bg set - b: {state['context_b']} - a: {state['context_a']}"
140

141
    tasks.add_task(bg, state)
142
    return state
143

144

145
# Sync versions
146

147

148
@app.get("/sync_async")
149
def get_sync_async(state: str = Depends(asyncgen_state)):
150
    return state
151

152

153
@app.get("/sync_sync")
154
def get_sync_sync(state: str = Depends(generator_state)):
155
    return state
156

157

158
@app.get("/sync_async_raise")
159
def get_sync_async_raise(state: str = Depends(asyncgen_state_try)):
160
    assert state == "asyncgen raise started"
161
    raise AsyncDependencyError()
162

163

164
@app.get("/sync_sync_raise")
165
def get_sync_sync_raise(state: str = Depends(generator_state_try)):
166
    assert state == "generator raise started"
167
    raise SyncDependencyError()
168

169

170
@app.get("/sync_async_raise_other")
171
def get_sync_async_raise_other(state: str = Depends(asyncgen_state_try)):
172
    assert state == "asyncgen raise started"
173
    raise OtherDependencyError()
174

175

176
@app.get("/sync_sync_raise_other")
177
def get_sync_sync_raise_other(state: str = Depends(generator_state_try)):
178
    assert state == "generator raise started"
179
    raise OtherDependencyError()
180

181

182
@app.get("/sync_context_b")
183
def get_sync_context_b(state: dict = Depends(context_b)):
184
    return state
185

186

187
@app.get("/sync_context_b_raise")
188
def get_sync_context_b_raise(state: dict = Depends(context_b)):
189
    assert state["context_b"] == "started b"
190
    assert state["context_a"] == "started a"
191
    raise OtherDependencyError()
192

193

194
@app.get("/sync_context_b_bg")
195
async def get_sync_context_b_bg(
196
    tasks: BackgroundTasks, state: dict = Depends(context_b)
197
):
198
    async def bg(state: dict):
199
        state[
200
            "sync_bg"
201
        ] = f"sync_bg set - b: {state['context_b']} - a: {state['context_a']}"
202

203
    tasks.add_task(bg, state)
204
    return state
205

206

207
@app.middleware("http")
208
async def middleware(request, call_next):
209
    response: StreamingResponse = await call_next(request)
210
    response.headers["x-state"] = json.dumps(state.copy())
211
    return response
212

213

214
client = TestClient(app)
215

216

217
def test_async_state():
218
    assert state["/async"] == "asyncgen not started"
219
    response = client.get("/async")
220
    assert response.status_code == 200, response.text
221
    assert response.json() == "asyncgen started"
222
    assert state["/async"] == "asyncgen completed"
223

224

225
def test_sync_state():
226
    assert state["/sync"] == "generator not started"
227
    response = client.get("/sync")
228
    assert response.status_code == 200, response.text
229
    assert response.json() == "generator started"
230
    assert state["/sync"] == "generator completed"
231

232

233
def test_async_raise_other():
234
    assert state["/async_raise"] == "asyncgen raise not started"
235
    with pytest.raises(OtherDependencyError):
236
        client.get("/async_raise_other")
237
    assert state["/async_raise"] == "asyncgen raise finalized"
238
    assert "/async_raise" not in errors
239

240

241
def test_sync_raise_other():
242
    assert state["/sync_raise"] == "generator raise not started"
243
    with pytest.raises(OtherDependencyError):
244
        client.get("/sync_raise_other")
245
    assert state["/sync_raise"] == "generator raise finalized"
246
    assert "/sync_raise" not in errors
247

248

249
def test_async_raise_raises():
250
    with pytest.raises(AsyncDependencyError):
251
        client.get("/async_raise")
252
    assert state["/async_raise"] == "asyncgen raise finalized"
253
    assert "/async_raise" in errors
254
    errors.clear()
255

256

257
def test_async_raise_server_error():
258
    client = TestClient(app, raise_server_exceptions=False)
259
    response = client.get("/async_raise")
260
    assert response.status_code == 500, response.text
261
    assert state["/async_raise"] == "asyncgen raise finalized"
262
    assert "/async_raise" in errors
263
    errors.clear()
264

265

266
def test_context_b():
267
    response = client.get("/context_b")
268
    data = response.json()
269
    assert data["context_b"] == "started b"
270
    assert data["context_a"] == "started a"
271
    assert state["context_b"] == "finished b with a: started a"
272
    assert state["context_a"] == "finished a"
273

274

275
def test_context_b_raise():
276
    with pytest.raises(OtherDependencyError):
277
        client.get("/context_b_raise")
278
    assert state["context_b"] == "finished b with a: started a"
279
    assert state["context_a"] == "finished a"
280

281

282
def test_background_tasks():
283
    response = client.get("/context_b_bg")
284
    data = response.json()
285
    assert data["context_b"] == "started b"
286
    assert data["context_a"] == "started a"
287
    assert data["bg"] == "not set"
288
    middleware_state = json.loads(response.headers["x-state"])
289
    assert middleware_state["context_b"] == "finished b with a: started a"
290
    assert middleware_state["context_a"] == "finished a"
291
    assert middleware_state["bg"] == "not set"
292
    assert state["context_b"] == "finished b with a: started a"
293
    assert state["context_a"] == "finished a"
294
    assert state["bg"] == "bg set - b: finished b with a: started a - a: finished a"
295

296

297
def test_sync_raise_raises():
298
    with pytest.raises(SyncDependencyError):
299
        client.get("/sync_raise")
300
    assert state["/sync_raise"] == "generator raise finalized"
301
    assert "/sync_raise" in errors
302
    errors.clear()
303

304

305
def test_sync_raise_server_error():
306
    client = TestClient(app, raise_server_exceptions=False)
307
    response = client.get("/sync_raise")
308
    assert response.status_code == 500, response.text
309
    assert state["/sync_raise"] == "generator raise finalized"
310
    assert "/sync_raise" in errors
311
    errors.clear()
312

313

314
def test_sync_async_state():
315
    response = client.get("/sync_async")
316
    assert response.status_code == 200, response.text
317
    assert response.json() == "asyncgen started"
318
    assert state["/async"] == "asyncgen completed"
319

320

321
def test_sync_sync_state():
322
    response = client.get("/sync_sync")
323
    assert response.status_code == 200, response.text
324
    assert response.json() == "generator started"
325
    assert state["/sync"] == "generator completed"
326

327

328
def test_sync_async_raise_other():
329
    with pytest.raises(OtherDependencyError):
330
        client.get("/sync_async_raise_other")
331
    assert state["/async_raise"] == "asyncgen raise finalized"
332
    assert "/async_raise" not in errors
333

334

335
def test_sync_sync_raise_other():
336
    with pytest.raises(OtherDependencyError):
337
        client.get("/sync_sync_raise_other")
338
    assert state["/sync_raise"] == "generator raise finalized"
339
    assert "/sync_raise" not in errors
340

341

342
def test_sync_async_raise_raises():
343
    with pytest.raises(AsyncDependencyError):
344
        client.get("/sync_async_raise")
345
    assert state["/async_raise"] == "asyncgen raise finalized"
346
    assert "/async_raise" in errors
347
    errors.clear()
348

349

350
def test_sync_async_raise_server_error():
351
    client = TestClient(app, raise_server_exceptions=False)
352
    response = client.get("/sync_async_raise")
353
    assert response.status_code == 500, response.text
354
    assert state["/async_raise"] == "asyncgen raise finalized"
355
    assert "/async_raise" in errors
356
    errors.clear()
357

358

359
def test_sync_sync_raise_raises():
360
    with pytest.raises(SyncDependencyError):
361
        client.get("/sync_sync_raise")
362
    assert state["/sync_raise"] == "generator raise finalized"
363
    assert "/sync_raise" in errors
364
    errors.clear()
365

366

367
def test_sync_sync_raise_server_error():
368
    client = TestClient(app, raise_server_exceptions=False)
369
    response = client.get("/sync_sync_raise")
370
    assert response.status_code == 500, response.text
371
    assert state["/sync_raise"] == "generator raise finalized"
372
    assert "/sync_raise" in errors
373
    errors.clear()
374

375

376
def test_sync_context_b():
377
    response = client.get("/sync_context_b")
378
    data = response.json()
379
    assert data["context_b"] == "started b"
380
    assert data["context_a"] == "started a"
381
    assert state["context_b"] == "finished b with a: started a"
382
    assert state["context_a"] == "finished a"
383

384

385
def test_sync_context_b_raise():
386
    with pytest.raises(OtherDependencyError):
387
        client.get("/sync_context_b_raise")
388
    assert state["context_b"] == "finished b with a: started a"
389
    assert state["context_a"] == "finished a"
390

391

392
def test_sync_background_tasks():
393
    response = client.get("/sync_context_b_bg")
394
    data = response.json()
395
    assert data["context_b"] == "started b"
396
    assert data["context_a"] == "started a"
397
    assert data["sync_bg"] == "not set"
398
    assert state["context_b"] == "finished b with a: started a"
399
    assert state["context_a"] == "finished a"
400
    assert (
401
        state["sync_bg"]
402
        == "sync_bg set - b: finished b with a: started a - a: finished a"
403
    )
404