在 FastAPI 中測試檔案上傳

FastAPI 是一款基於 Starlette 的 Web 框架,其在 API 的開發體驗令人驚豔。

註:FastAPI 官方宣稱,得益於 Starlette 及 Pydantic,它的效能甚至能夠與 Go 比肩;然而根據 TechEmpower Web Framework Benchmarks 在 2022 年 7 月 19 日的測試,其效能約為 278 及 279 名,事實上以效能而言它位於中下游的水準。

與 Laravel 這類包山包海的全能型框架不同,FastAPI 選擇了一條「微」框架的道路,它更像是 Gin(Go) 的設計:輕量、精簡,並且在有需要時讓開發者自行安裝。這種設計很大程度給予開發者自由,甚至連資料夾結構都沒有官方預設(如果願意的話,甚至可以只靠一個 app.pymain.py 就建構起整個 API 服務)。

然而最讓我感到驚豔的,非屬自動 API 文件生成功能。眾所周知,工程師是種「最討厭別人不寫文件,但又不喜歡自己寫文件」的生物,FastAPI 內建由 Pydantic Model 進行自動化文件生成的功能,這能夠很大程度上減少撰寫文件的工作量。

實作

檔案上傳功能

實作一個用戶頭像(Avatar)上傳功能,本次主要聚焦在「上傳功能」,對於一些比較不重要的細節會加以省略:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
# main.py
from fastapi import FastAPI, UploadFile, Depends, status, HTTPException

from schemas import User

app = FastAPI()

@app.post("/profile/avatar")
def upload_avatar(
    avatar: UploadFile,
    user: User = Depends(get_authed_user) # get_authed_user 用來取得當前登入的用戶,實作細節省略
):
    if avatar.content_type not in ("imgage/jpg", "image/jpeg", "image/png"):
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="avatar should be an image")

    user.update_avatar(avatar) # 實作如何處理當前上傳的 avatar,可以放在 local 或上傳到 S3

    return {"message": "avatar upload successfully"}

測試上傳功能

為了方便測試,通常會先建立 conftest.py 用於設定 pytest

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
# conftest.py
import pytest
from typing import Generator
from fastapi.testclient import TestClient

from main import app
from main.schemas import User

def get_testing_user():
    return User(email="[email protected]", password="password")

@pytest.fixture(scope="module")
def client() -> Generator:
    # 覆寫 get_authed_user,使所有在 TestClient 中的請求都可以使用假的 User 當成已登入用戶
    app.dependency_overrides[get_authed_user] = lambda: User(email="[email protected]", password="password")

    with TestClient(app) as c:
        yield c

測試上傳功能

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import tempfile
from fastapi.testclient import TestClient

def test_upload_avatar(client: TestClient):
    with tempfile.NamedTemporaryFile(suffix=".png") as f:
        response = client.post(
            "/profile/avatar",
            # 檔案上傳應該提供一個 dict[str, list],並且 list 中應該分別是 (檔案名, 檔案資源, 檔案 MIME)
            files={"avatar": ("foobar.png", f, "image/png")},
        )

        assert response.status_code == 200

利用 Python 內建的 tempfile 暫時性地開啟檔案,並且在使用完畢後自動回收。

同場加映

上傳至 s3

 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
# config.py

from pydantic import BaseSettings
from boto3.session import Session

class Settings(BaseSettings):
    AWS_ACCESS_KEY_ID: str | None = None
    AWS_SECRET_ACCESS_KEY: str | None = None
    AWS_S3_ENDPOINT_URL: str | None = None
    AWS_S3_BUCKET: str | None = None

    class Config:
        env_file = ".env"
        case_sensitive = True


settins = Settings()

s3_session = Session(
    aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
    aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
)

def get_s3_resource():
    return s3_session.resource("s3", endpoint_url=settings.AWS_S3_ENDPOINT_URL).Bucket(settings.AWS_S3_BUCKET)

在預設上,雖然 boto3 可以直接讀取 ~/.aws/credentials~/.aws/config,理論上不必特別設定。

然而,為了統一 config 的設定方式,這邊將其抽出來另外存取,這麼一來也可以配合其它服務(如 DB)或 SDK 去做設定。

 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
# main.py
@app.post("/profile/avatar")
def upload_avatar(
    avatar: UploadFile,
    user: User = Depends(get_authed_user), # get_authed_user 用來取得當前登入的用戶,實作細節省略
    s3 = Depends(get_s3_resource)
):
    if avatar.content_type not in ("imgage/jpg", "image/jpeg", "image/png"):
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="avatar should be an image")

    obj_name = handle_uploaded_file(avatar, user.id, s3) # 建立檔名,並且上傳至 S3
    user.update_avatar(obj_name) # 將檔名更新至 user 的 avatar 參數中

    return {"message": "avatar upload successfully"}

def handle_uploaded_file(file: UploadFile, user_id: int, s3) -> str:
    obj_name = "avatar/{}-{}{}".format(
        # 以 user_id 作為前綴
        user_id, 
        # 加入 5 個隨機小寫字母字元
        ''.join(random.choices(string.ascii_lowercase, k=5)),
        # 保持相同的副檔名
        os.path.splitext(avatar.filename)[1]
    )

    s3.upload_fileobj(file.file, obj_name)

    return obj_name