使用 pytest 在 Flask 中测试服务器发送的事件

Question

我有一个 Flask 应用程序，我使用服务器发送的事件将数据发送到我的前端。

@bp.route("/stream", methods=("GET",))
def stream_translations():
    translation_schema = TranslationSchema()

    def event_stream():
        while True:
            recently_updated = [
                translation_schema.dump(translation)
                for translation in recently_updated_translations()]
            if recently_updated:
                yield f"data: {json.dumps(recently_updated)}\n\n"

    return Response(event_stream(), mimetype="text/event-stream")

它工作得很好，但我也想为其编写一个测试来确定。我以前从未为生成器编写过测试，并且绝对不是服务器发送的事件。目前，这就是我所拥有的：

def test_stream(client):
    response = client.get("/translations/stream")

    assert response.status_code == 200
    assert response.mimetype == "text/event-stream"

当然这只是测试响应，但我还想测试event_stream()生成器。我该怎么做呢？

Answer 1

通过一些小的重构，我们可以通过仔细的测试模拟成功测试我们的 SSE 端点。

步骤 1. 为无限循环添加“旁路”

测试此代码时您可能面临的第一个问题是 SSE 事件流提倡使用while True循环。这在浏览器环境中有意义，但在服务器端集成测试中意义不大，因为它会导致您的测试用例“挂起”。

我们可以通过将生成器代码重构为辅助函数来解决此问题：

import time

def event_stream(timeout = 0.0):
    starting_time = time.time()
    while not timeout or (time.time() - starting_time < timeout):
        ...
        yield f'data:{json.dumps(...)}'

这样，您可以重构原始stream_translations函数，如下所示：

@bp.route("/stream", methods=("GET",))
def stream_translations():
    return Response(event_stream(), mimetype='text/event-stream')

通过这次重构，我们完成了两件事：

1. event_stream有一个附加timeout参数，它允许我们防止测试中的无限循环

2. event_stream不是嵌套函数，这允许我们适当地模拟它。

步骤 2. 绕过测试中的无限循环

理想情况下，我们只想timeout在测试中指定参数，但允许它在生产设置中永远继续。我们可以通过模拟来实现这一点：

from contextlib import contextmanager
from functools import partial
from unittest import mock

@contextmanager
def mock_events():
    with mock.patch(
        # NOTE: For example, if your function is found in app/views/translations.py,
        # then this import path would be 'app.views.translations.event_stream'
        'python.import.path.to.stream_translations.event_stream',

        # NOTE: Remember to import this function wherever you defined it.
        partial(event_stream, timeout=1.5),
    ):
        yield

这使我们能够最小化测试代码和生产代码之间的差异（因为运行相同的函数），但实际上允许该函数在测试中完成。

使用此函数，您的测试可能如下所示：

def test_stream(client):
    with mock_events():
        response = client.get("/translations/stream")

    assert response.data.decode()