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?
With some minor refactoring, we can successfully test our SSE endpoints through careful test mocks.
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:
event_stream
has an additional timeout
parameter, which allows us to prevent the infinite loop in tests
event_stream
is not a nested function, which allows us to mock it appropriately.
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()
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
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.