FastAPI测试中SqlAlchemy查询在WebSocket场景下挂起问题排查

云上老周

2026-8-2

问题描述

我开发了一个带WebSocket的FastAPI应用并进行测试，需通过异步测试在测试过程中校验数据库数据。测试环境配置如下：

数据库会话Fixture定义

@pytest_asyncio.fixture(scope="session")
async def async_engine() -> AsyncEngine:
    yield create_async_engine(settings.TEST_DATABASE_URL)

@pytest_asyncio.fixture(scope="function")
async def session(async_engine):
    AsyncSessionLocal = sessionmaker(
        bind=async_engine, autoflush=False, expire_on_commit=False, future=True, class_=AsyncSession
    )
    yield AsyncSessionLocal()

客户端Fixture定义

@pytest_asyncio.fixture(scope="function")
async def client(session):
    app.dependency_overrides[get_db_session] = lambda: session
    async with AsyncClient(transport=ASGIWebSocketTransport(app), base_url="http://test") as client:
        yield client

其中transport=ASGIWebSocketTransport(app)来自httpx-ws库，原生httpx不支持WebSocket。

测试用例

@pytest.mark.asyncio
async def test_http(session, client):
    r = await client.get('/test/')

@pytest.mark.asyncio
async def test_ws(session, client):
    async with aconnect_ws("/ws/", client) as ws:
        await ws.send_json({})

test_http执行正常，但test_ws出现问题：/ws/接口中所有SQL查询均永久挂起。当引擎开启echo=True时，日志中能看到类似sqlalchemy.engine.Engine SELECT test.id FROM test的语句，但后续无任何输出。我已尝试查看pg_stat_activity、pg_locks和pg_blocking_pids(pid)但未找到原因。目前我已通读httpx-ws代码未发现潜在问题，请问还能通过哪些方式排查数据库连接挂起的原因？

排查建议

检查事件循环一致性
异步SQLAlchemy依赖稳定的事件循环上下文。WebSocket处理时可能存在事件循环切换，或代码中混入了同步阻塞操作（比如误用time.sleep()而非asyncio.sleep()），导致数据库查询无法获得调度时间。可以在数据库查询前后添加print(asyncio.get_event_loop())，对比HTTP测试与WebSocket测试的事件循环是否一致。
验证会话的上下文管理
检查/ws/接口中数据库会话的使用方式，是否正确通过async with session.begin():开启事务上下文。若会话未正确进入事务，查询可能卡在等待状态。同时确认WebSocket路由中是否存在未正确释放的会话引用，导致连接池耗尽。
开启SQLAlchemy详细调试日志
除echo=True外，配置SQLAlchemy的连接池和引擎日志为DEBUG级别，追踪连接的获取、释放过程：
```
import logging
logging.basicConfig()
logging.getLogger('sqlalchemy.pool').setLevel(logging.DEBUG)
logging.getLogger('sqlalchemy.engine').setLevel(logging.DEBUG)
```
查看是否存在连接被占用未释放、连接池为空的情况。
编写极简WebSocket测试示例
实现一个仅执行SELECT 1;的极简WebSocket路由，排除业务逻辑干扰。若该示例仍挂起，说明问题出在会话与WebSocket的集成层面；若正常，则逐步添加业务代码定位问题点。

修复会话Fixture的清理逻辑
当前session fixture未做事务回滚和会话清理，HTTP测试后会话可能处于异常状态（如未提交事务），影响WebSocket测试。修改fixture添加清理步骤：

@pytest_asyncio.fixture(scope="function")
async def session(async_engine):
    AsyncSessionLocal = sessionmaker(
        bind=async_engine, autoflush=False, expire_on_commit=False, future=True, class_=AsyncSession
    )
    async with AsyncSessionLocal() as session:
        async with session.begin():
            yield session
        await session.rollback()

也可以为WebSocket测试单独创建独立的会话fixture，避免与HTTP测试共享实例。

确认异步数据库驱动配置
确保使用异步PostgreSQL驱动（如asyncpg），且连接URL格式正确（postgresql+asyncpg://user:pass@host/db）。误用同步驱动会直接导致异步查询挂起，同时检查asyncpg与SQLAlchemy的版本兼容性。

追踪协程状态与栈信息
为测试添加超时逻辑，超时后打印所有活跃协程的状态和调用栈，定位阻塞点：

import asyncio
from contextlib import suppress

@pytest.mark.asyncio
async def test_ws(session, client):
    try:
        async with asyncio.timeout(10):
            async with aconnect_ws("/ws/", client) as ws:
                await ws.send_json({})
    except asyncio.TimeoutError:
        print("活跃协程状态：")
        for task in asyncio.all_tasks():
            print(f"协程 {task.get_name()}: {task._state}")
            with suppress(Exception):
                task.print_stack()

内容的提问来源于stack exchange，提问作者Artem Shalagin

FastAPI测试中SqlAlchemy查询在WebSocket场景下挂起问题排查

数据库会话Fixture定义 #

客户端Fixture定义 #

测试用例 #

数据库会话Fixture定义

客户端Fixture定义

测试用例