mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-28 10:54:19 +00:00
0b6cd45950
# What does this PR do? This expands the `test_sse` test suite and fixes some edge cases with bugs in our SSE error handling to ensure streaming clients always get a proper error response. First, we handle the case where a client disconnects before we actually start streaming the response back. Previously we only handled the case where a client disconnected as we were streaming the response, but there was an edge case where a client disconnecting before we streamed any response back did not trigger our logic to cleanly handle that disconnect. Second, we handle the case where an error is thrown from the server before the actual async generator gets created from the provider. This happens in scenarios like the newly merged OpenAI API input validation, where we eagerly raise validation errors before returning the async generator object that streams the responses back. ## Test Plan Tested via: ``` python -m pytest -s -v tests/unit/server/test_sse.py ``` Both test cases failed before, and passed afterwards. The test cases were written based on me experimenting with actual clients that would do bad things like randomly disconnect or send invalid input in streaming mode and I hit these two cases, where things were misbehaving in our error handling. Signed-off-by: Ben Browning <bbrownin@redhat.com>
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from llama_stack.distribution.server.server import create_sse_event, sse_generator
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_generator_basic():
|
|
# An AsyncIterator wrapped in an Awaitable, just like our web methods
|
|
async def async_event_gen():
|
|
async def event_gen():
|
|
yield "Test event 1"
|
|
yield "Test event 2"
|
|
|
|
return event_gen()
|
|
|
|
sse_gen = sse_generator(async_event_gen())
|
|
assert sse_gen is not None
|
|
|
|
# Test that the events are streamed correctly
|
|
seen_events = []
|
|
async for event in sse_gen:
|
|
seen_events.append(event)
|
|
assert len(seen_events) == 2
|
|
assert seen_events[0] == create_sse_event("Test event 1")
|
|
assert seen_events[1] == create_sse_event("Test event 2")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_generator_client_disconnected():
|
|
# An AsyncIterator wrapped in an Awaitable, just like our web methods
|
|
async def async_event_gen():
|
|
async def event_gen():
|
|
yield "Test event 1"
|
|
# Simulate a client disconnect before emitting event 2
|
|
raise asyncio.CancelledError()
|
|
|
|
return event_gen()
|
|
|
|
sse_gen = sse_generator(async_event_gen())
|
|
assert sse_gen is not None
|
|
|
|
seen_events = []
|
|
async for event in sse_gen:
|
|
seen_events.append(event)
|
|
|
|
# We should see 1 event before the client disconnected
|
|
assert len(seen_events) == 1
|
|
assert seen_events[0] == create_sse_event("Test event 1")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_generator_client_disconnected_before_response_starts():
|
|
# Disconnect before the response starts
|
|
async def async_event_gen():
|
|
raise asyncio.CancelledError()
|
|
|
|
sse_gen = sse_generator(async_event_gen())
|
|
assert sse_gen is not None
|
|
|
|
seen_events = []
|
|
async for event in sse_gen:
|
|
seen_events.append(event)
|
|
|
|
# No events should be seen since the client disconnected immediately
|
|
assert len(seen_events) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_generator_error_before_response_starts():
|
|
# Raise an error before the response starts
|
|
async def async_event_gen():
|
|
raise Exception("Test error")
|
|
|
|
sse_gen = sse_generator(async_event_gen())
|
|
assert sse_gen is not None
|
|
|
|
seen_events = []
|
|
async for event in sse_gen:
|
|
seen_events.append(event)
|
|
|
|
# We should have 1 error event
|
|
assert len(seen_events) == 1
|
|
assert 'data: {"error":' in seen_events[0]
|