chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,281 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user