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

Question

I have a Flask application where I use server-sent events to send data to my front-end.

@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")

It works fine, but I also want to write a test for it to make sure. I've never written a test for a generator before, and definitely not server-sent events. Currently this is what I have:

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

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

Of course this just tests the response, but I also want to test the event_stream() generator. How do I do this?

Answer 1

With some minor refactoring, we can successfully test our SSE endpoints through careful test mocks.

Step 1. Adding a "Bypass" for the Infinite Loop

The first issue that you may be facing when testing this code is that SSE event streams promotes using while True loops. This makes sense in a browser environment, but not so much in a server-side integration test, since it will cause your test cases to "hang".

We can address this by refactoring your generator code into a helper function:

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(...)}'

With this, you can refactor your original stream_translations function as such:

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

With this refactor, we accomplish two things:

1. event_stream has an additional timeout parameter, which allows us to prevent the infinite loop in tests

2. event_stream is not a nested function, which allows us to mock it appropriately.

Step 2. Bypassing the Infinite Loop in Tests

Ideally, we want to only specify the timeout parameter in tests, but allow it to continue forever in production settings. We can accomplish this through 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

This allows us to minimize differences between test code and production code (since the same function is run), but actually allows this function to complete in tests.

With this function, your test may look something like this:

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

    assert response.data.decode()

Answer 2

After a while my friend solve it very easy using 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