如何使用pytest测试基于FastAPI实现的SSE（Server-Sent Events）端点？

问题场景 #

基于FastAPI实现的SSE端点可正常运行，但使用TestClient编写pytest测试时，代码在client.get('/')处停滞，无法执行后续断言逻辑。

问题原因 #

TestClient默认会等待整个响应流完全结束后才返回结果，但SSE是无限持续的数据流，端点中的生成器会一直循环发送事件，导致测试进程被无限阻塞。

解决方案 #

方案1：开启流式请求并限制读取行数 #

在发送请求时指定stream=True，让客户端立即返回响应对象，读取指定数量的事件后主动终止循环并关闭响应，避免无限等待。

修改后的测试代码：

# test_main.py
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from fastapi.responses import StreamingResponse
import asyncio
import json   

app = FastAPI()

@app.get('/')
async def main(request: Request) -> StreamingResponse:
    async def gen():
        while True: 
            yield 'data: ' + json.dumps({'msg': 'Hello World!'}) + '\n\n'
            await asyncio.sleep(0.5)
            
    return StreamingResponse(gen(), media_type='text/event-stream')
    
client = TestClient(app)
    
def test_read_stream():
    # 开启流式请求，不等待响应完全结束
    with client.get("/", stream=True) as response:
        assert response.status_code == 200
        # 读取前2条有效事件后终止
        count = 0
        for line in response.iter_lines():
            if line:  # 跳过SSE中的空行
                assert line == 'data: {"msg": "Hello World!"}'
                count += 1
                if count >= 2:
                    break

方案2：给SSE生成器添加终止逻辑 #

在端点的生成器中增加检测逻辑，比如判断客户端是否断开连接，或者通过查询参数控制事件发送次数，让生成器在测试场景下主动结束流。

修改后的端点代码：

@app.get('/')
async def main(request: Request) -> StreamingResponse:
    async def gen():
        try:
            send_count = 0
            # 通过查询参数控制测试时的最大发送次数
            max_count = request.query_params.get("max_count", None)
            
            while True: 
                yield 'data: ' + json.dumps({'msg': 'Hello World!'}) + '\n\n'
                await asyncio.sleep(0.5)
                
                # 测试场景下，达到指定次数后停止
                if max_count and send_count >= int(max_count):
                    break
                send_count += 1
                
                # 检测客户端是否断开连接，真实场景下自动终止
                if await request.is_disconnected():
                    break
        except asyncio.CancelledError:
            pass
            
    return StreamingResponse(gen(), media_type='text/event-stream')

对应的测试代码：

def test_read_stream():
    # 传递参数控制只发送2条事件
    response = client.get("/?max_count=2")
    assert response.status_code == 200
    # 过滤空行后验证事件数量和内容
    valid_lines = [line for line in response.iter_lines() if line]
    assert len(valid_lines) == 2
    for line in valid_lines:
        assert line == 'data: {"msg": "Hello World!"}'

说明 #

方案1无需修改端点代码，适合快速验证SSE格式和内容；方案2更贴近真实场景的连接管理，同时支持测试时的可控终止。

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