[Day 25] 好用的測試模組：Pytest

發表於 更新於 分類於

今天來聊聊簡單的主題 ── 測試

API 測試

這邊我們就直接來看看怎麼測試 FastAPI 的 API。

首先，要先安裝 pytest 和 httpx

1
pip install pytest httpx

使用 pytest 的原因很簡單，就是讓我們可以執行測試的套件，而 httpx 則是讓我們使用 TestClient，它就像是另一個使用者，它會向我們的 FastAPI 後端發送請求，之後再利用 assert 判斷結果是否符合我們的預期。

這邊直接來看範例，這是簡單的主程式

1
2
3
4
5
6
7
8
# main.py
from fastapi import FastAPI

app = FastAPI()

@app.get("/")
async def read_main():
    return {"msg": "Hello World"}

而這是測試腳本

1
2
3
4
5
6
7
8
9
10
11
# test_main.py

from fastapi.testclient import TestClient
from main import app

client = TestClient(app)

def test_read_main():
    response = client.get("/")
    assert response.status_code == 200
    assert response.json() == {"msg": "Hello World"}

需要注意的是，測試腳本和裡面的 function 都必須使用 test 開頭，才會被 pytest 納入執行範圍。

接著在 terminal 輸入

1
pytest test_main.py

稍等一下，就可以看到結果了

如果想要同時測試所有腳本，則可以省去後面指定的檔案路徑

1
pytest

測試啟動/關閉事件的測試

除了 API，FastAPI 還有很多設定，而這些東西基本上也都可以寫進測試中，例如：啟動/關閉事件。

什麼是啟動、關閉事件

之前沒有機會介紹，所以只好偷偷放在這邊了

在 FastAPI 中，我們可以在它啟動或關閉時幫我們做一些事，例如：紀錄 log、載入機器學習模型，讓後面的 API 不用一直反覆的讀取這樣的一個大檔案。

寫法也很容易

1
2
3
4
5
6
7
8
# main.py
@app.on_event("startup")
async def startup_event():
    print("System Start")

@app.on_event("shutdown")
def shutdown():
    print("System Stop")

而這個事件的測試可以這樣寫

1
2
3
4
5
6
# test_main.py
def test_read_items():
    with TestClient(app) as client:
        response = client.get("/items/foo")
        assert response.status_code == 200
        assert response.json() == {"name": "Fighters"}
1
2
3
4
@app.on_event("startup")
async def startup_event():
    items["foo"] = {"name": "Fighters"}
    items["bar"] = {"name": "Tenders"}

最後額外補充一下，在今年 3 月，FastAPI 的 0.93.0 版多了一個新功能 lifespan，目的是取代原本的 startup、shutdown 事件，這邊直接看一下官方範例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from contextlib import asynccontextmanager
from fastapi import FastAPI

def fake_answer_to_everything_ml_model(x: float):
    return x * 42

ml_models = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the ML model
    ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model
    yield
    # Clean up the ML models and release the resources
    ml_models.clear()

app = FastAPI(lifespan=lifespan)

@app.get("/predict")
async def predict(x: float):
    result = ml_models["answer_to_everything"](x)
    return {"result": result}

可以看到，它在 lifespan 內用 yield 把兩個階段整合在一起，這樣的好處是，過去在 startup、shutdown 事件中無法共用的變數就變成可以共用了。

覆蓋原本設定 (Depends、middleware)

Depends

在某些時候，為了方便測試，我們希望把原本寫在 API 內的 Depends() 給取消或覆蓋掉，例如：登入驗證。

與其在測試時先跑一遍登入 API 拿到 JWT，再放入 header 再進行測試，不如直接把原本的驗證移除，如果有需要使用者資訊，那可以改成用一個有固定回傳的函數來覆蓋掉原本的登入驗證。

當然，在某些時候，能完整這樣測試會更好

這個做法就是 dependency_overrides，可以來看看官方範例，首先是 API 的部分

1
2
3
4
5
6
7
8
async def common_parameters(
    q: Union[str, None] = None, skip: int = 0, limit: int = 100
):
    return {"q": q, "skip": skip, "limit": limit}

@app.get("/items/")
async def read_items(commons: Annotated[dict, Depends(common_parameters)]):
    return {"message": "Hello Items!", "params": commons}

測試的時候，就用 app.dependency_overrides()，讓 override_dependency() 覆蓋原本的 common_parameters

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from fastapi.testclient import TestClient
from typing import Annotated, Union
from main import app

client = TestClient(app)
async def override_dependency(q: Union[str, None] = None):
    return {"q": q, "skip": 5, "limit": 10}
app.dependency_overrides[common_parameters] = override_dependency

def test_override_in_items():
    response = client.get("/items/")
    assert response.status_code == 200
    assert response.json() == {
        "message": "Hello Items!",
        "params": {"q": None, "skip": 5, "limit": 10},
    }

middleware

同樣道理，或許有些時候這個 Depends 會放在 middleware，或是再 middleware 有做什麼特殊設計，導致在測試時不好用，此時就可以考慮直接移除 middleware

1
2
3
4
5
6
from fastapi.testclient import TestClient
from main import app

client = TestClient(app)
app.user_middleware.clear()
app.middleware_stack = app.build_middleware_stack()

小結

今天我們簡單的介紹了怎麼寫 FastAPI 的測試，不過，測試的領域水也不淺，更重要的是怎麼設計測試腳本，這邊就沒辦法多做介紹了。