wehub-resource-sync
60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import deque
|
|
|
|
import pytest
|
|
|
|
from agents.sandbox.session.pty_output import collect_pty_output
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_pty_output_waits_for_notification() -> None:
|
|
output_chunks: deque[bytes] = deque()
|
|
output_lock = asyncio.Lock()
|
|
output_notify = asyncio.Event()
|
|
done = False
|
|
|
|
async def produce_output() -> None:
|
|
nonlocal done
|
|
await asyncio.sleep(0)
|
|
async with output_lock:
|
|
output_chunks.append(b"notified output")
|
|
done = True
|
|
output_notify.set()
|
|
|
|
producer_task = asyncio.create_task(produce_output())
|
|
output, original_token_count = await collect_pty_output(
|
|
output_chunks=output_chunks,
|
|
output_lock=output_lock,
|
|
output_notify=output_notify,
|
|
is_done=lambda: done,
|
|
yield_time_ms=500,
|
|
max_output_tokens=None,
|
|
)
|
|
await producer_task
|
|
|
|
assert output == b"notified output"
|
|
assert original_token_count is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_pty_output_drains_chunks_added_when_done() -> None:
|
|
output_chunks = deque([b"before done"])
|
|
|
|
def mark_done() -> bool:
|
|
output_chunks.append(b" after done")
|
|
return True
|
|
|
|
output, original_token_count = await collect_pty_output(
|
|
output_chunks=output_chunks,
|
|
output_lock=asyncio.Lock(),
|
|
output_notify=asyncio.Event(),
|
|
is_done=mark_done,
|
|
yield_time_ms=500,
|
|
max_output_tokens=None,
|
|
)
|
|
|
|
assert output == b"before done after done"
|
|
assert original_token_count is None
|