feat: add OCI GenAI streaming agent foundation
This commit is contained in:
319
tests/test_oci_stream.py
Normal file
319
tests/test_oci_stream.py
Normal file
@@ -0,0 +1,319 @@
|
||||
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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user