使用 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 参数，但允许它在生产设置中永远持续下去。我们可以通过 mocks 来实现：

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()

Answer 2

过了一会儿，我的朋友使用pytest 很容易解决它：

@fixture(scope='function')
def stream_start(client: FlaskClient):
    client.post("/stream/start")
    sleep(4)  # you get data for 4 sec
    client.post("/stream/stop")

def test_output_stream(client: FlaskClient, api_start):  
    with client.get('/stream/output_stream') as res:
        for encoded in res.iter_encoded():
            if encoded:
                data = json.loads(encoded.decode('UTF-8')[::])
                assert real_data == data