FastAPI与PostgreSQL单元测试隔离方案咨询

我开发了一个包含多个端点与模型的FastAPI接口，为提升后续开发者工作效率计划编写自动化测试，但遇到测试隔离问题：

• 测试数据会直接写入生产数据库，风险极高
• 尝试创建独立测试数据库，但模型直接绑定了生产环境的特定finance schema
• 尝试用Docker搭建测试数据库，但异步数据库连接配置失败

附上单元测试代码与模型代码，求可行的测试隔离方案。

单元测试代码 #

import pytest
from sqlalchemy.orm import sessionmaker
from database.conn import Base
from app.models.finance.bills_to_pay_models import BillsToPay  
import os
from sqlalchemy.future import select
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

USER = os.getenv('DB_USER')
PASSWORD = os.getenv('DB_PASSWORD')
HOST = os.getenv('DB_HOST')
NAME = os.getenv('DB_NAME_TEST_OMIE')

SQLALCHEMY_DATABASE_URL = f"postgresql+asyncpg://{USER}:{PASSWORD}@{HOST}/{NAME}"

engine = create_async_engine(SQLALCHEMY_DATABASE_URL, echo=True)
SessionLocal = sessionmaker(bind=engine, class_=AsyncSession)

@pytest.fixture(scope="function")
async def test_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    async with SessionLocal() as session:
        transaction = await session.begin()
        yield
        await transaction.rollback()
        await conn.run_sync(Base.metadata.drop_all)  
        await session.close()

@pytest.mark.asyncio
async def test_bills_to_pay_create(test_db):
    async with SessionLocal() as session:
        new_category = BillsToPay(
            example1="Teste",
            example2=12,
            example3=1,
        )

        try:
            session.add(new_category)
            await session.commit()

        except Exception as e:
            await session.rollback()
            raise e

        result = await session.execute(select(BillsToPay).filter_by(example1="Teste"))
        category_from_db = result.first()[0]
        if not category_from_db:
            raise ValueError("Error")

        assert category_from_db.example1 == "Teste"
        assert category_from_db.example2 == 12
        assert category_from_db.example3 == 1
        
        await session.close()

模型代码 #

from database.conn import Base
from sqlalchemy import String, Integer, Column 

class BillsToPay(Base):
    __tablename__  = 'BillsToPay'
    __table_args__ = {'schema': 'finance'}

    id = Column(Integer, primary_key=True, index=True)
    example1 = Column(String, nullable=True)
    example2 = Column(Integer, nullable=True)
    example3 = Column(Integer, nullable=True)

1. 动态切换模型Schema，解除生产绑定 #

模型硬编码schema='finance'是核心限制，改成通过环境变量动态设置，不影响生产代码，测试时自动切换到测试schema（比如test_finance）。

修改模型代码： #

from database.conn import Base
from sqlalchemy import String, Integer, Column 
import os

# 从环境变量取schema，默认用生产的finance
DB_SCHEMA = os.getenv("DB_SCHEMA", "finance")

class BillsToPay(Base):
    __tablename__  = 'BillsToPay'
    __table_args__ = {'schema': DB_SCHEMA}

    id = Column(Integer, primary_key=True, index=True)
    example1 = Column(String, nullable=True)
    example2 = Column(Integer, nullable=True)
    example3 = Column(Integer, nullable=True)

测试时设置环境变量： #

在测试代码开头或者pytest配置文件（conftest.py）里添加：

import os
os.environ["DB_SCHEMA"] = "test_finance"

这样测试时模型会自动指向测试schema，生产环境保持默认的finance。

2. 优化异步测试的事务隔离，避免重复建表 #

你当前的测试fixture存在上下文管理问题：conn在engine.begin()的上下文结束后已关闭，yield后调用conn.run_sync会报错。改成嵌套事务+回滚的方式，不用每次删表，测试效率更高。

修复后的pytest fixture： #

import pytest
import os
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from database.conn import Base

# 测试数据库配置
USER = os.getenv('DB_USER')
PASSWORD = os.getenv('DB_PASSWORD')
HOST = os.getenv('DB_HOST')
NAME = os.getenv('DB_NAME_TEST_OMIE')
SQLALCHEMY_DATABASE_URL = f"postgresql+asyncpg://{USER}:{PASSWORD}@{HOST}/{NAME}"

# 全局测试引擎和会话工厂
test_engine = create_async_engine(SQLALCHEMY_DATABASE_URL, echo=True)
TestSessionLocal = sessionmaker(bind=test_engine, class_=AsyncSession, expire_on_commit=False)

@pytest.fixture(scope="session")
async def setup_test_db():
    # 会话级别：创建所有表，测试结束后删除
    async with test_engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield
    async with test_engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)

@pytest.fixture(scope="function")
async def db_session(setup_test_db):
    # 函数级别：开启嵌套事务，测试后回滚
    session = TestSessionLocal()
    await session.begin_nested()
    
    # 覆盖commit方法，实际用flush避免提交到数据库
    original_commit = session.commit
    def mock_commit():
        return session.flush()
    session.commit = mock_commit

    yield session

    # 回滚嵌套事务，清理数据
    await session.rollback()
    await session.close()

修改测试用例： #

@pytest.mark.asyncio
async def test_bills_to_pay_create(db_session):
    new_category = BillsToPay(
        example1="Teste",
        example2=12,
        example3=1,
    )

    db_session.add(new_category)
    await db_session.commit()  # 实际是flush，不会提交到数据库

    result = await db_session.execute(select(BillsToPay).filter_by(example1="Teste"))
    category_from_db = result.scalar_one_or_none()
    
    assert category_from_db is not None
    assert category_from_db.example1 == "Teste"
    assert category_from_db.example2 == 12
    assert category_from_db.example3 == 1

这个方案的优势：

• 会话级别只建一次表，测试速度大幅提升
• 函数级别用嵌套事务回滚，每个测试数据完全隔离
• 避免频繁删表的IO开销

3. Docker Compose管理异步测试数据库 #

之前Docker失败大概率是配置问题，用Docker Compose一键启动PostgreSQL异步数据库，测试时自动连接。

编写docker-compose.test.yml： #

version: '3.8'
services:
  test-db:
    image: postgres:15-alpine
    environment:
      POSTGRES_USER: test_user
      POSTGRES_PASSWORD: test_pass
      POSTGRES_DB: test_omie_db
    ports:
      - "5433:5432"
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U test_user -d test_omie_db"]
      interval: 5s
      timeout: 5s
      retries: 5

启动测试数据库： #

docker-compose -f docker-compose.test.yml up -d

设置测试环境变量： #

export DB_USER=test_user
export DB_PASSWORD=test_pass
export DB_HOST=localhost
export DB_NAME_TEST_OMIE=test_omie_db
export DB_SCHEMA=test_finance

测试结束后关闭容器：

docker-compose -f docker-compose.test.yml down

额外建议 #

• 把测试配置放到单独的.env.test文件，用python-dotenv加载，避免手动设置环境变量
• 在pytest配置文件conftest.py里统一管理测试数据库的初始化和清理逻辑
• 测试时禁用SQLAlchemy的echo模式，减少日志输出干扰

内容的提问来源于stack exchange，提问作者Ivan Campelo