wehub-resource-sync
282 lines
8.4 KiB
Python
282 lines
8.4 KiB
Python
import asyncio
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from e2b.envd.process import process_pb2
|
|
from e2b.sandbox_async.commands.command_handle import AsyncCommandHandle
|
|
from e2b.sandbox_sync.commands.command_handle import CommandHandle
|
|
|
|
EMOJI = "😀"
|
|
EMOJI_BYTES = EMOJI.encode("utf-8") # 4 bytes
|
|
|
|
|
|
def _stdout_event(data: bytes) -> process_pb2.StartResponse:
|
|
return process_pb2.StartResponse(
|
|
event=process_pb2.ProcessEvent(
|
|
data=process_pb2.ProcessEvent.DataEvent(stdout=data)
|
|
)
|
|
)
|
|
|
|
|
|
def _stderr_event(data: bytes) -> process_pb2.StartResponse:
|
|
return process_pb2.StartResponse(
|
|
event=process_pb2.ProcessEvent(
|
|
data=process_pb2.ProcessEvent.DataEvent(stderr=data)
|
|
)
|
|
)
|
|
|
|
|
|
def _end_event(exit_code: int = 0) -> process_pb2.StartResponse:
|
|
return process_pb2.StartResponse(
|
|
event=process_pb2.ProcessEvent(
|
|
end=process_pb2.ProcessEvent.EndEvent(
|
|
exit_code=exit_code, exited=True, status="exited"
|
|
)
|
|
)
|
|
)
|
|
|
|
|
|
async def _kill() -> bool:
|
|
return True
|
|
|
|
|
|
class _AsyncControllableEvents:
|
|
"""Async event source that delivers items on demand.
|
|
|
|
Lets a test hold the handle's event-handling task blocked waiting for the
|
|
next event (idle between bursts), then push a late event after
|
|
``disconnect()`` to confirm it never reaches the callback — the transport
|
|
condition that triggers the JS SDK leak.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._queue: asyncio.Queue = asyncio.Queue()
|
|
self._closed = False
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
if self._closed and self._queue.empty():
|
|
raise StopAsyncIteration
|
|
item = await self._queue.get()
|
|
if item is None:
|
|
raise StopAsyncIteration
|
|
return item
|
|
|
|
def push(self, event):
|
|
self._queue.put_nowait(event)
|
|
|
|
async def aclose(self):
|
|
# Intentionally does NOT unblock a read that is already suspended on
|
|
# the queue: this models the JS-leak transport condition where the
|
|
# stream stays open past the close request. The only thing that stops
|
|
# the handle's event loop here is task cancellation in disconnect().
|
|
self._closed = True
|
|
|
|
|
|
async def test_async_disconnect_stops_callbacks():
|
|
events = _AsyncControllableEvents()
|
|
chunks = []
|
|
handle = AsyncCommandHandle(
|
|
pid=1,
|
|
handle_kill=_kill,
|
|
# The handle only async-iterates and aclose()s the stream; this stand-in
|
|
# satisfies both without being a real async generator.
|
|
events=cast(Any, events),
|
|
on_stdout=chunks.append,
|
|
)
|
|
|
|
# First burst is delivered to the live subscriber.
|
|
events.push(_stdout_event(b"a"))
|
|
for _ in range(200):
|
|
if chunks == ["a"]:
|
|
break
|
|
await asyncio.sleep(0.005)
|
|
assert chunks == ["a"]
|
|
|
|
# disconnect() cancels the event-handling task and closes the stream, so
|
|
# once it returns the callback must not fire again.
|
|
await handle.disconnect()
|
|
|
|
# A late event (stdout arriving after disconnect) must never reach the
|
|
# callback.
|
|
events.push(_stdout_event(b"b"))
|
|
await asyncio.sleep(0.05)
|
|
assert chunks == ["a"]
|
|
|
|
|
|
def test_sync_has_no_background_subscription():
|
|
# The sync handle has no detached event-handling task: events are consumed
|
|
# only while the caller iterates (e.g. inside wait()), so there is no
|
|
# subscription that could keep firing after disconnect().
|
|
consumed = []
|
|
|
|
def events():
|
|
consumed.append("started")
|
|
yield _stdout_event(b"a")
|
|
yield _end_event()
|
|
|
|
handle = CommandHandle(pid=1, handle_kill=lambda: True, events=events())
|
|
|
|
# Nothing is consumed until the caller iterates.
|
|
assert consumed == []
|
|
|
|
# disconnect() just closes the (un-started) stream — still nothing consumed.
|
|
handle.disconnect()
|
|
assert consumed == []
|
|
|
|
|
|
def test_sync_records_result_before_yielding_flushed_chunk():
|
|
# A consumer that stops iterating right after the end event's flushed chunk
|
|
# must still observe the exit code: the result is recorded before the
|
|
# flushed chunk is yielded.
|
|
def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
yield _end_event(0)
|
|
|
|
handle = CommandHandle(pid=1, handle_kill=lambda: True, events=events())
|
|
iterator = iter(handle)
|
|
assert next(iterator) == ("a", None, None)
|
|
# The end event flushes a trailing replacement character; pull just that
|
|
# chunk and then stop iterating.
|
|
next(iterator)
|
|
iterator.close()
|
|
|
|
assert handle._result is not None
|
|
assert handle._result.exit_code == 0
|
|
assert handle._result.stdout == "a�"
|
|
|
|
|
|
def test_sync_decodes_multibyte_chars_split_across_chunks():
|
|
def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
yield _stdout_event(EMOJI_BYTES[2:] + b"b")
|
|
yield _stderr_event(EMOJI_BYTES[:3])
|
|
yield _stderr_event(EMOJI_BYTES[3:])
|
|
yield _end_event()
|
|
|
|
chunks = []
|
|
handle = CommandHandle(pid=1, handle_kill=lambda: True, events=events())
|
|
result = handle.wait(on_stdout=chunks.append)
|
|
|
|
assert result.stdout == f"a{EMOJI}b"
|
|
assert result.stderr == EMOJI
|
|
assert "�" not in result.stdout
|
|
assert "�" not in result.stderr
|
|
assert "".join(chunks) == f"a{EMOJI}b"
|
|
|
|
|
|
def test_sync_replaces_incomplete_trailing_utf8():
|
|
def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
yield _end_event()
|
|
|
|
handle = CommandHandle(pid=1, handle_kill=lambda: True, events=events())
|
|
result = handle.wait()
|
|
|
|
assert result.stdout == "a�"
|
|
|
|
|
|
async def test_async_decodes_multibyte_chars_split_across_chunks():
|
|
async def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
yield _stdout_event(EMOJI_BYTES[2:] + b"b")
|
|
yield _stderr_event(EMOJI_BYTES[:3])
|
|
yield _stderr_event(EMOJI_BYTES[3:])
|
|
yield _end_event()
|
|
|
|
chunks = []
|
|
handle = AsyncCommandHandle(
|
|
pid=1,
|
|
handle_kill=_kill,
|
|
events=events(),
|
|
on_stdout=chunks.append,
|
|
)
|
|
result = await handle.wait()
|
|
|
|
assert result.stdout == f"a{EMOJI}b"
|
|
assert result.stderr == EMOJI
|
|
assert "�" not in result.stdout
|
|
assert "�" not in result.stderr
|
|
assert "".join(chunks) == f"a{EMOJI}b"
|
|
|
|
|
|
async def test_async_replaces_incomplete_trailing_utf8():
|
|
async def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
yield _end_event()
|
|
|
|
handle = AsyncCommandHandle(pid=1, handle_kill=_kill, events=events())
|
|
result = await handle.wait()
|
|
|
|
assert result.stdout == "a�"
|
|
|
|
|
|
def test_sync_flushes_incomplete_trailing_utf8_without_end_event():
|
|
def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
|
|
chunks = []
|
|
handle = CommandHandle(pid=1, handle_kill=lambda: True, events=events())
|
|
for stdout, _, _ in handle:
|
|
if stdout is not None:
|
|
chunks.append(stdout)
|
|
|
|
assert "".join(chunks) == "a�"
|
|
|
|
|
|
async def test_async_flushes_incomplete_trailing_utf8_without_end_event():
|
|
async def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
|
|
chunks = []
|
|
handle = AsyncCommandHandle(
|
|
pid=1,
|
|
handle_kill=_kill,
|
|
events=events(),
|
|
on_stdout=chunks.append,
|
|
)
|
|
await handle._wait
|
|
|
|
assert "".join(chunks) == "a�"
|
|
|
|
|
|
def test_sync_flushes_incomplete_trailing_utf8_on_stream_error():
|
|
def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
raise RuntimeError("stream died")
|
|
|
|
chunks = []
|
|
handle = CommandHandle(pid=1, handle_kill=lambda: True, events=events())
|
|
|
|
# The stream raises before an end event, but the buffered bytes must still
|
|
# be flushed as a replacement character before the error is surfaced.
|
|
with pytest.raises(RuntimeError):
|
|
for stdout, _, _ in handle:
|
|
if stdout is not None:
|
|
chunks.append(stdout)
|
|
|
|
assert "".join(chunks) == "a�"
|
|
|
|
|
|
async def test_async_flushes_incomplete_trailing_utf8_on_stream_error():
|
|
async def events():
|
|
yield _stdout_event(b"a" + EMOJI_BYTES[:2])
|
|
raise RuntimeError("stream died")
|
|
|
|
chunks = []
|
|
handle = AsyncCommandHandle(
|
|
pid=1,
|
|
handle_kill=_kill,
|
|
events=events(),
|
|
on_stdout=chunks.append,
|
|
)
|
|
await handle._wait
|
|
|
|
# The stream raised before an end event, but the buffered bytes must still
|
|
# be flushed to the stdout callback as a replacement character.
|
|
assert "".join(chunks) == "a�"
|
|
assert isinstance(handle._iteration_exception, RuntimeError)
|