62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
"""Shared helpers for the Omnigent REPL pexpect tests.
|
||
|
||
These helpers are used by multiple ``test_repl_*.py`` files in
|
||
this directory. They are intentionally kept separate from
|
||
``_pexpect_harness.py`` (the phase 0 foundation) to avoid
|
||
diverging the foundation file across worktrees during the
|
||
phase 0 merge — the foundation is shared with another agent's
|
||
work-in-progress and must remain byte-identical there.
|
||
|
||
If a helper here proves stable across phases it can later be
|
||
folded into ``_pexpect_harness.py`` as part of a deliberate
|
||
foundation-update commit.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import pexpect
|
||
|
||
# Per-chunk read window for :func:`drain_for`. Tuned so that
|
||
# prompt-toolkit's default ~50 ms refresh tick still fires at
|
||
# least once inside one chunk — keeps draining responsive
|
||
# without wasting a full chunk-budget on each idle slice.
|
||
_CHUNK_TIMEOUT_SEC = 0.3
|
||
|
||
|
||
def drain_for(child: pexpect.spawn, total_timeout: float) -> str:
|
||
"""
|
||
Collect all PTY output that arrives within a bounded
|
||
window.
|
||
|
||
Issues repeated short ``read_nonblocking`` calls and
|
||
concatenates the chunks. More robust than a single long
|
||
``read_nonblocking`` when prompt-toolkit emits a render
|
||
burst across multiple frames — a single read can return
|
||
after the first chunk and miss follow-up frames painted
|
||
by the next refresh tick.
|
||
|
||
The function returns early if the PTY idles for one
|
||
chunk-interval before the budget elapses, so the wall-
|
||
clock cost is bounded by whichever is shorter.
|
||
|
||
:param child: Live ``pexpect.spawn`` child returned by
|
||
:func:`tests.e2e.omnigent._pexpect_harness.spawn_omnigent_run`.
|
||
:param total_timeout: Overall seconds to spend draining,
|
||
e.g. ``5.0``. A typical value for "let the post-key
|
||
render frames settle" is 2–6 seconds.
|
||
:returns: Concatenated raw PTY output collected during the
|
||
drain window. May be empty if the PTY produced no
|
||
output during the budget.
|
||
"""
|
||
chunks: list[str] = []
|
||
remaining = total_timeout
|
||
while remaining > 0:
|
||
chunk_timeout = min(_CHUNK_TIMEOUT_SEC, remaining)
|
||
try:
|
||
chunk = child.read_nonblocking(size=65536, timeout=chunk_timeout)
|
||
except pexpect.TIMEOUT:
|
||
break
|
||
chunks.append(chunk)
|
||
remaining -= chunk_timeout
|
||
return "".join(chunks)
|