320 lines
7.5 KiB
Python
320 lines
7.5 KiB
Python
import pytest
|
|
import asyncio
|
|
import threading
|
|
|
|
from types import SimpleNamespace
|
|
from contextlib import aclosing
|
|
|
|
from voice_assistant.agent.oci_stream import stream_sse_events
|
|
|
|
|
|
class FakeSSEStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
|
|
def events(self):
|
|
yield SimpleNamespace(data='{"text":"Olá"}')
|
|
yield SimpleNamespace(data='{"text":" mundo"}')
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stream_sse_events_delivers_data():
|
|
stream = FakeSSEStream()
|
|
|
|
received = []
|
|
|
|
async for data in stream_sse_events(stream):
|
|
received.append(data)
|
|
|
|
assert received == [
|
|
'{"text":"Olá"}',
|
|
'{"text":" mundo"}',
|
|
]
|
|
|
|
assert stream.closed is True
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stream_sse_events_propagates_worker_error():
|
|
class FailingSSEStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
|
|
def events(self):
|
|
yield SimpleNamespace(data='{"text":"Olá"}')
|
|
raise RuntimeError("Simulated OCI SSE read failure.")
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
stream = FailingSSEStream()
|
|
received = []
|
|
|
|
with pytest.raises(
|
|
RuntimeError,
|
|
match="Simulated OCI SSE read failure",
|
|
):
|
|
async for data in stream_sse_events(stream):
|
|
received.append(data)
|
|
|
|
assert received == ['{"text":"Olá"}']
|
|
assert stream.closed is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stream_sse_events_cancellation_closes_stream():
|
|
"""
|
|
Verify that cancelling the async consumer closes the SSE stream
|
|
and stops its worker.
|
|
"""
|
|
class BlockingSSEStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
self.started = threading.Event()
|
|
self.release = threading.Event()
|
|
|
|
def events(self):
|
|
self.started.set()
|
|
|
|
# Simulate a blocking synchronous HTTP read.
|
|
self.release.wait(timeout=5.0)
|
|
|
|
if not self.closed:
|
|
yield SimpleNamespace(
|
|
data='{"text":"Late event"}'
|
|
)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
self.release.set()
|
|
|
|
stream = BlockingSSEStream()
|
|
|
|
async def consume():
|
|
async for _ in stream_sse_events(
|
|
stream,
|
|
shutdown_timeout=1.0,
|
|
):
|
|
pass
|
|
|
|
consumer = asyncio.create_task(consume())
|
|
|
|
try:
|
|
started = await asyncio.to_thread(
|
|
stream.started.wait,
|
|
1.0,
|
|
)
|
|
|
|
assert started, "SSE worker did not start."
|
|
|
|
consumer.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await asyncio.wait_for(
|
|
consumer,
|
|
timeout=2.0,
|
|
)
|
|
|
|
assert stream.closed is True
|
|
|
|
finally:
|
|
stream.release.set()
|
|
|
|
if not consumer.done():
|
|
consumer.cancel()
|
|
|
|
await asyncio.gather(
|
|
consumer,
|
|
return_exceptions=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stream_sse_events_applies_backpressure():
|
|
"""
|
|
Verify that the SSE worker stops advancing when
|
|
the asynchronous consumer is slower than the producer.
|
|
"""
|
|
class FastSSEStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
self.produced = 0
|
|
self.lock = threading.Lock()
|
|
|
|
def events(self):
|
|
for index in range(100):
|
|
with self.lock:
|
|
self.produced += 1
|
|
|
|
yield SimpleNamespace(
|
|
data=f'{{"index":{index}}}'
|
|
)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
stream = FastSSEStream()
|
|
|
|
received = []
|
|
|
|
async for data in stream_sse_events(
|
|
stream,
|
|
max_queue_size=2,
|
|
):
|
|
received.append(data)
|
|
|
|
# Slow down the consumer to exercise backpressure.
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert len(received) == 100
|
|
assert stream.produced == 100
|
|
assert stream.closed is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stream_sse_events_blocks_producer_when_queue_full():
|
|
"""
|
|
Verify that the synchronous producer cannot advance
|
|
indefinitely while the async consumer is paused.
|
|
"""
|
|
class CountingSSEStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
self.produced = 0
|
|
self.lock = threading.Lock()
|
|
|
|
def events(self):
|
|
for index in range(100):
|
|
with self.lock:
|
|
self.produced += 1
|
|
|
|
yield SimpleNamespace(
|
|
data=f'{{"index":{index}}}'
|
|
)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
stream = CountingSSEStream()
|
|
|
|
consumer_paused = asyncio.Event()
|
|
release_consumer = asyncio.Event()
|
|
|
|
async def consume():
|
|
received = []
|
|
|
|
async for data in stream_sse_events(
|
|
stream,
|
|
max_queue_size=2,
|
|
):
|
|
received.append(data)
|
|
|
|
if len(received) == 1:
|
|
consumer_paused.set()
|
|
await release_consumer.wait()
|
|
|
|
return received
|
|
|
|
task = asyncio.create_task(consume())
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
consumer_paused.wait(),
|
|
timeout=1.0,
|
|
)
|
|
|
|
# Allow the worker thread to reach the queue capacity.
|
|
await asyncio.sleep(0.05)
|
|
|
|
with stream.lock:
|
|
produced_while_paused = stream.produced
|
|
|
|
# One consumed event, two buffered events, and at most
|
|
# one event already fetched by the worker.
|
|
assert produced_while_paused <= 4
|
|
|
|
finally:
|
|
release_consumer.set()
|
|
|
|
if not task.done():
|
|
try:
|
|
await asyncio.wait_for(task, timeout=2.0)
|
|
except TimeoutError:
|
|
task.cancel()
|
|
await asyncio.gather(
|
|
task,
|
|
return_exceptions=True,
|
|
)
|
|
|
|
assert stream.closed is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_stream_sse_events_cancellation_with_full_queue():
|
|
"""
|
|
Verify that cancelling a paused consumer with a full queue
|
|
does not leave the SSE worker waiting indefinitely.
|
|
"""
|
|
class FastSSEStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
self.produced = 0
|
|
self.lock = threading.Lock()
|
|
|
|
def events(self):
|
|
for index in range(100):
|
|
with self.lock:
|
|
self.produced += 1
|
|
|
|
yield SimpleNamespace(
|
|
data=f'{{"index":{index}}}'
|
|
)
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
stream = FastSSEStream()
|
|
consumer_paused = asyncio.Event()
|
|
|
|
async def consume():
|
|
async with aclosing(
|
|
stream_sse_events(
|
|
stream,
|
|
max_queue_size=1,
|
|
shutdown_timeout=1.0,
|
|
)
|
|
) as events:
|
|
async for _ in events:
|
|
consumer_paused.set()
|
|
await asyncio.Event().wait()
|
|
|
|
task = asyncio.create_task(consume())
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
consumer_paused.wait(),
|
|
timeout=1.0,
|
|
)
|
|
|
|
task.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await asyncio.wait_for(
|
|
task,
|
|
timeout=2.0,
|
|
)
|
|
|
|
assert stream.closed is True
|
|
|
|
finally:
|
|
if not task.done():
|
|
task.cancel()
|
|
|
|
await asyncio.gather(
|
|
task,
|
|
return_exceptions=True,
|
|
)
|
|
|