133 lines
2.6 KiB
Python
133 lines
2.6 KiB
Python
import asyncio
|
|
import threading
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from voice_assistant.agent.oci_connection import open_oci_stream
|
|
|
|
|
|
class FakeStream:
|
|
def __init__(self):
|
|
self.closed = False
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_open_oci_stream_success():
|
|
stream = FakeStream()
|
|
|
|
response = SimpleNamespace(data=stream)
|
|
|
|
result = await open_oci_stream(lambda: response)
|
|
|
|
assert result is response
|
|
assert stream.closed is False
|
|
|
|
stream.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_open_oci_stream_closes_late_response():
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
finished = threading.Event()
|
|
|
|
stream = FakeStream()
|
|
response = SimpleNamespace(data=stream)
|
|
|
|
def open_connection():
|
|
started.set()
|
|
release.wait(timeout=5.0)
|
|
|
|
try:
|
|
return response
|
|
finally:
|
|
finished.set()
|
|
|
|
task = asyncio.create_task(
|
|
open_oci_stream(open_connection)
|
|
)
|
|
|
|
try:
|
|
assert await asyncio.to_thread(
|
|
started.wait,
|
|
1.0,
|
|
)
|
|
|
|
task.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
finally:
|
|
release.set()
|
|
|
|
assert await asyncio.to_thread(
|
|
finished.wait,
|
|
1.0,
|
|
)
|
|
|
|
# The worker may still be completing its cleanup.
|
|
for _ in range(100):
|
|
if stream.closed:
|
|
break
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert stream.closed is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_open_oci_stream_cancellation_during_handoff():
|
|
"""
|
|
Verify that a response is not abandoned when cancellation
|
|
happens around completion of the opening operation.
|
|
"""
|
|
stream = FakeStream()
|
|
response = SimpleNamespace(data=stream)
|
|
|
|
opened = threading.Event()
|
|
release = threading.Event()
|
|
|
|
def open_connection():
|
|
opened.set()
|
|
release.wait(timeout=5.0)
|
|
return response
|
|
|
|
task = asyncio.create_task(
|
|
open_oci_stream(open_connection)
|
|
)
|
|
|
|
try:
|
|
assert await asyncio.to_thread(
|
|
opened.wait,
|
|
1.0,
|
|
)
|
|
|
|
release.set()
|
|
task.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
for _ in range(100):
|
|
if stream.closed:
|
|
break
|
|
await asyncio.sleep(0.001)
|
|
|
|
assert stream.closed is True
|
|
|
|
finally:
|
|
release.set()
|
|
|
|
if not task.done():
|
|
task.cancel()
|
|
|
|
await asyncio.gather(
|
|
task,
|
|
return_exceptions=True,
|
|
)
|
|
|