本文檔詳細解釋了在 pytest 環(huán)境下，如何使用 SQLAlchemy 的異步功能來管理數(shù)據(jù)庫連接和狀態(tài)，以確保測試的隔離性和可靠性。

1. Pytest Fixture 核心概念：作用域 (Scope)

@pytest.fixture 是 pytest 中一個非常強大的功能，它用于為測試函數(shù)提供數(shù)據(jù)、對象或預(yù)備/清理環(huán)境。其中，scope 參數(shù)是控制 fixture生命周期的關(guān)鍵。

scope 參數(shù)決定了一個 fixture 實例會被創(chuàng)建和銷毀的頻率。它有以下幾個選項，按從小到大的順序排列：

在我們的數(shù)據(jù)庫測試策略中，我們組合使用了不同的 scope：

• db_engine 使用 scope="module": 因為數(shù)據(jù)庫引擎的創(chuàng)建開銷較大，我們希望在一個測試文件中它只被創(chuàng)建一次。
• db_session 使用 scope="function": 因為我們希望每個測試函數(shù)都在一個獨立的、干凈的事務(wù)中運行，測試結(jié)束后立即回滾，互不干擾。

2. 核心概念：MetaData 對象

MetaData 對象可以被看作是您數(shù)據(jù)庫 schema 在 Python 代碼中的一個“注冊表”或“目錄”。

它是一個容器，用于存放所有與它關(guān)聯(lián)的 Table 對象的定義。

在我們的項目中，src/models/tables.py 中定義的所有 Table 對象都注冊到了一個全局的 metadata 實例上。

這使得我們可以執(zhí)行強大的 schema 級別的操作，如 metadata.create_all() 和 metadata.drop_all()。

2. 測試隔離策略

為了確保每個單元測試都在一個獨立、干凈的環(huán)境中運行，不受其他測試的影響，我們需要實現(xiàn)一種“測試隔離”策略。主要有兩種方法：

策略一：每次重建 (Recreation per Test) - 我們當前使用的方法

這是最直觀、最健壯的方法，不依賴特定數(shù)據(jù)庫的事務(wù)特性。

示例代碼

# test/dao/test_user_dao.py
@pytest.mark.asyncio
async def test_create_user():
    # 1. 為此測試創(chuàng)建獨立的引擎
    engine = create_async_engine(DATABASE_URL)
    
    # 2. 在測試開始時，物理刪除并重建所有表
    async with engine.begin() as conn:
        await conn.run_sync(metadata.drop_all)
        await conn.run_sync(metadata.create_all)

    # 3. 執(zhí)行測試邏輯 (包括 commit)
    async with AsyncSession(engine) as session:
        # ... DAO 調(diào)用 ...
    
    # 4. 在測試結(jié)束時，銷毀引擎以關(guān)閉所有物理連接
    await engine.dispose()

關(guān)鍵點解析

metadata.drop_all() / create_all():

• 作用: 實現(xiàn)測試隔離。
• 原理: 在每個測試函數(shù)開始時，物理地刪除并重新創(chuàng)建所有數(shù)據(jù)庫表。這保證了每個測試面對的都是一個全新的、空的數(shù)據(jù)庫。
• DDL 與事務(wù): 您提出了一個很好的問題：DROP/CREATE 是 DDL，為何要放在 engine.begin() 事務(wù)塊中？在 PostgreSQL 中，DDL 是事務(wù)性的，可以被包含在事務(wù)中。但在這里，使用 engine.begin() 的主要目的是為了優(yōu)雅地管理連接的生命周期（獲取連接、執(zhí)行操作、釋放連接），而不是為了 DDL 的原子性。

engine.dispose():

• 作用: 實現(xiàn)資源清理。
• 原理: 關(guān)閉并銷毀 engine 內(nèi)部維護的整個連接池中的所有物理數(shù)據(jù)庫連接。

與 rollback 的區(qū)別: dispose() 不會回滾任何已提交的事務(wù)。它只負責(zé)關(guān)閉網(wǎng)絡(luò)連接。在我們的例子中，測試數(shù)據(jù)的清理是由下一個測試開始時的 drop_all 完成的。

• 優(yōu)點: 極其可靠，跨數(shù)據(jù)庫兼容性好。

• 缺點: 性能較低。對于每個測試都刪除和創(chuàng)建表，開銷很大。

策略二：事務(wù)回滾 (Transaction Rollback) - 更高效的策略

這是一種更高級、性能更好的方法，它利用了數(shù)據(jù)庫的事務(wù)特性。

理論代碼

# conftest.py - (這是一個理論上的例子，我們當前項目沒有使用)

@pytest.fixture(scope="session")
async def engine():
    # 整個測試會話只創(chuàng)建一個引擎
    db_engine = create_async_engine(DATABASE_URL)
    yield db_engine
    await db_engine.dispose()

@pytest.fixture(scope="session", autouse=True)
async def setup_database(engine):
    # 在會話開始時創(chuàng)建一次表，結(jié)束時刪除一次
    async with engine.begin() as conn:
        await conn.run_sync(metadata.create_all)
    yield
    async with engine.begin() as conn:
        await conn.run_sync(metadata.drop_all)

@pytest.fixture(scope="function")
async def db_session(engine) -> AsyncSession:
    # 為每個測試函數(shù)提供一個特殊的“回滾”會話
    async with engine.connect() as connection:
        async with connection.begin() as transaction: # 開始一個事務(wù)
            async with AsyncSession(bind=connection) as session:
                yield session
                # 測試結(jié)束后，回滾這個事務(wù)，撤銷所有 DML 操作
                await transaction.rollback()

關(guān)鍵點解析

setup_database Fixture: 在整個測試會話開始時創(chuàng)建一次所有表，在會話結(jié)束時刪除它們。

db_session Fixture:

• connection.begin(): 在每個測試函數(shù)開始時，它會啟動一個事務(wù)（或者在支持的數(shù)據(jù)庫上是一個嵌套事務(wù)/保存點）。
• yield session: 測試函數(shù)在自己的這個“子事務(wù)”中運行，可以自由地 COMMIT 數(shù)據(jù)。
• await transaction.rollback(): 這是核心。當測試函數(shù)結(jié)束時，無論測試成功與否，也無論函數(shù)內(nèi)部是否執(zhí)行了 commit，這個 fixture 都會強制回滾最外層的事務(wù)。

效果: test_create_user 中 COMMIT 的數(shù)據(jù)實際上只被提交到了一個未關(guān)閉的事務(wù)中。測試一結(jié)束，整個事務(wù)就被回滾，數(shù)據(jù)庫瞬間恢復(fù)到測試開始前的狀態(tài)。

• 優(yōu)點: 速度極快。ROLLBACK 是一個非常輕量級的操作。
• 缺點: 實現(xiàn)更復(fù)雜，且依賴于數(shù)據(jù)庫對事務(wù)性 DDL 的支持。

總結(jié)

我們當前采用的**策略一（每次重建）**雖然性能稍低，但它更簡單、直觀，并且能 100% 保證每個測試的隔離性。對于大多數(shù)項目來說，這都是一個非常可靠和推薦的起點。

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論