62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
|
|
import pexpect
|
|
|
|
|
|
CURSOR_POSITION_QUERY = "\x1b[6n"
|
|
COMMAND_DONE = "__BRUSH_E2E_COMMAND_DONE__"
|
|
# Seconds to wait for any one expected string. The overall wait is bounded separately, since a
|
|
# shell that keeps asking where the cursor is would otherwise refresh this on every reply.
|
|
TIMEOUT = int(os.environ.get("E2E_TEST_TIMEOUT", "10"))
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
def spawn(command, args, env, prompt):
|
|
child = pexpect.spawn(
|
|
command,
|
|
args,
|
|
dimensions=(30, 100),
|
|
encoding="utf-8",
|
|
env=env | {"BRUSH_E2E_COMMAND_DONE": COMMAND_DONE},
|
|
timeout=TIMEOUT,
|
|
)
|
|
LOG.info("shell startup:\n%s", expect_output(child, prompt))
|
|
return child
|
|
|
|
|
|
def expect_output(shell, expected):
|
|
output = []
|
|
deadline = time.monotonic() + TIMEOUT
|
|
while True:
|
|
try:
|
|
matched = shell.expect_exact([expected, CURSOR_POSITION_QUERY])
|
|
except pexpect.exceptions.ExceptionPexpect:
|
|
output.append(shell.before)
|
|
LOG.exception("PTY error waiting for %r:\n%s", expected, "".join(output))
|
|
raise
|
|
output.append(shell.before)
|
|
if matched == 0:
|
|
return "".join(output)
|
|
if time.monotonic() >= deadline:
|
|
LOG.error("PTY gave up waiting for %r:\n%s", expected, "".join(output))
|
|
raise pexpect.exceptions.TIMEOUT(
|
|
f"still answering cursor-position queries after {TIMEOUT}s waiting for {expected!r}"
|
|
)
|
|
# Claim the cursor is at the top left. Some prompts (starship) query it at every redraw
|
|
# and wait for the reply before printing anything.
|
|
shell.send("\x1b[1;1R")
|
|
|
|
|
|
def run(shell, prompt, command):
|
|
shell.sendline(
|
|
f"{command}; __brush_e2e_status=$?; "
|
|
'printf "\\n%s\\n" "$BRUSH_E2E_COMMAND_DONE"; '
|
|
'(exit "$__brush_e2e_status")'
|
|
)
|
|
output = expect_output(shell, COMMAND_DONE)
|
|
output += expect_output(shell, prompt)
|
|
LOG.info("$ %s\n%s", command, output)
|
|
return output
|