在 FastAPI 中測試檔案上傳

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

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

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

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

實作

檔案上傳功能

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

# 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

# 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

測試上傳功能

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

# 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 去做設定。

# 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