290 lines
9.8 KiB
Python
290 lines
9.8 KiB
Python
# Licensed under the Apache License, Version 2.0 (the "License"); you may
|
|
# not use this file except in compliance with the License. You may obtain
|
|
# a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
# License for the specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
import contextlib
|
|
import fcntl
|
|
import itertools
|
|
import multiprocessing
|
|
import os
|
|
import pty
|
|
import re
|
|
import signal
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
import termios
|
|
import time
|
|
import traceback
|
|
|
|
import types
|
|
from typing import Optional, Generator
|
|
import typing
|
|
|
|
|
|
(LINES, COLUMNS) = (24, 80)
|
|
|
|
|
|
def _open_fifo(path: Optional[str],
|
|
write: bool) -> Optional[int]:
|
|
if path is None:
|
|
return None
|
|
return os.open(path, os.O_WRONLY if write else os.O_RDONLY)
|
|
|
|
|
|
class IsolationEnvironment:
|
|
def __init__(self, pid: int, ptm_fd: int,
|
|
stdin_fifo: Optional[str] = None,
|
|
stdout_fifo: Optional[str] = None,
|
|
stderr_fifo: Optional[str] = None):
|
|
self._pid = pid
|
|
|
|
self._ptm_fd = ptm_fd
|
|
self._tty = os.fdopen(ptm_fd, 'r')
|
|
|
|
self._stdin_fifo_fd = _open_fifo(stdin_fifo, True)
|
|
self._stdout_fifo_fd = _open_fifo(stdout_fifo, False)
|
|
self._stderr_fifo_fd = _open_fifo(stderr_fifo, False)
|
|
|
|
self._exit_code: Optional[int] = None
|
|
|
|
def interrupt(self) -> None:
|
|
self.write(b'\x03')
|
|
time.sleep(0.001)
|
|
|
|
def write(self, data: bytes) -> None:
|
|
os.write(self._ptm_fd, data)
|
|
|
|
def readline(self) -> str:
|
|
return self._tty.readline()
|
|
|
|
def error_output(self) -> str:
|
|
if self._stderr_fifo_fd is None:
|
|
return ''
|
|
with os.fdopen(self._stderr_fifo_fd) as errf:
|
|
self._stderr_fifo_fd = None
|
|
return errf.read()
|
|
|
|
def stdout_pipe(self) -> typing.TextIO:
|
|
assert self._stdout_fifo_fd is not None
|
|
stdout = os.fdopen(self._stdout_fifo_fd)
|
|
self._stdout_fifo_fd = None
|
|
return stdout
|
|
|
|
def stdin_pipe(self) -> typing.TextIO:
|
|
assert self._stdin_fifo_fd is not None
|
|
return os.fdopen(self._stdin_fifo_fd, 'w', closefd=False)
|
|
|
|
def close(self, get_return_code: typing.Callable[[], int]) -> None:
|
|
for i in range(100):
|
|
if os.waitpid(self._pid, os.WNOHANG) != (0, 0):
|
|
break
|
|
time.sleep(0.001)
|
|
else:
|
|
os.kill(self._pid, signal.SIGTERM)
|
|
os.waitpid(self._pid, 0)
|
|
self._exit_code = get_return_code()
|
|
self._tty.close()
|
|
for fifo_fd in (self._stdin_fifo_fd,
|
|
self._stdout_fifo_fd,
|
|
self._stderr_fifo_fd):
|
|
if fifo_fd is not None:
|
|
os.close(fifo_fd)
|
|
|
|
def exit_code(self) -> int:
|
|
assert self._exit_code is not None
|
|
return self._exit_code
|
|
|
|
|
|
class PagerControl:
|
|
_page_end = None
|
|
_ctrl_chars = re.compile(r'\x1b\['
|
|
r'(\?|[0-2]?K|[0-9]*;?[0-9]*H|[0-3]?J|[0-9]*m)')
|
|
|
|
def __init__(self, isolation_env: IsolationEnvironment):
|
|
self.env = isolation_env
|
|
self._total_lines = 0
|
|
self._lines = self._iter_lines()
|
|
|
|
def _iter_lines(self) -> Generator[typing.Union[str, None],
|
|
None, None]:
|
|
while True:
|
|
line = '\x1b[?'
|
|
while line.lstrip(' q').startswith('\x1b[?'):
|
|
rawline = self.env.readline()
|
|
line = (rawline.replace('\x07', '') # Ignore bell
|
|
.replace('\x1b[m', '')) # Ignore style reset
|
|
before, reset, after = line.partition('\x1b[2J')
|
|
for segment in filter(bool, (before, after)):
|
|
if segment == '\x1b[1m~\x1b[0m\n':
|
|
# Ignore lines that are filling blank vertical space with
|
|
# '~' after hitting Ctrl-C when the screen is not full
|
|
continue
|
|
visible = self._ctrl_chars.sub('', segment)
|
|
if ((visible.rstrip() == ':' or '(END)' in visible)
|
|
and segment.replace('\x1b[m', '') != visible):
|
|
yield self._page_end
|
|
elif visible != '\n' or segment == visible:
|
|
self._total_lines += 1
|
|
yield visible
|
|
|
|
def read_lines(self, count: int) -> typing.Iterator[str]:
|
|
return itertools.islice((line for line in self._lines
|
|
if line is not self._page_end), count)
|
|
|
|
def _lines_in_page(self) -> int:
|
|
original_count = self._total_lines
|
|
try:
|
|
for line in self._lines:
|
|
if line is self._page_end:
|
|
break
|
|
except IOError:
|
|
pass
|
|
return self._total_lines - original_count
|
|
|
|
def advance(self) -> int:
|
|
self.env.write(b' ')
|
|
return self._lines_in_page()
|
|
|
|
def quit(self) -> int:
|
|
self.env.write(b'q')
|
|
return self._lines_in_page()
|
|
|
|
def total_lines(self) -> int:
|
|
return self._total_lines
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _fifo(fifo_path: str) -> Generator[str, None, None]:
|
|
try:
|
|
os.mkfifo(fifo_path, 0o600)
|
|
yield fifo_path
|
|
finally:
|
|
try:
|
|
os.unlink(fifo_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _fifos(*fifo_names: Optional[str]) -> Generator[typing.List[Optional[str]],
|
|
None, None]:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with contextlib.ExitStack() as stack:
|
|
fifos = [stack.enter_context(_fifo(os.path.join(directory,
|
|
name)))
|
|
if name is not None else None
|
|
for name in fifo_names]
|
|
yield fifos
|
|
|
|
|
|
class TestProcessNotComplete(Exception):
|
|
pass
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def isolate(child_function: typing.Callable[[], int],
|
|
stdin_pipe: bool = False,
|
|
stdout_pipe: bool = False,
|
|
stderr_pipe: bool = True,
|
|
*,
|
|
lines: int = LINES,
|
|
columns: int = COLUMNS) -> Generator[IsolationEnvironment,
|
|
None, None]:
|
|
with _fifos('stdin' if stdin_pipe else None,
|
|
'stdout' if stdout_pipe else None,
|
|
'stderr' if stderr_pipe else None) as fifo_paths:
|
|
result_r, result_w = os.pipe()
|
|
env_pid, tty = pty.fork()
|
|
if env_pid == 0:
|
|
try:
|
|
os.close(result_r)
|
|
pts = os.ttyname(sys.stdout.fileno())
|
|
|
|
ctx = multiprocessing.get_context('spawn')
|
|
p = ctx.Process(target=_run_test, args=(child_function, pts,
|
|
*fifo_paths))
|
|
p.start()
|
|
|
|
def handle_terminate(signum: int,
|
|
frame: types.FrameType) -> None:
|
|
if p.is_alive():
|
|
p.terminate()
|
|
|
|
signal.signal(signal.SIGTERM, handle_terminate)
|
|
# Signals from the pty are for the test process, not us
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
|
|
p.join(2) # Wait at most 2s
|
|
except BaseException:
|
|
with os.fdopen(result_w, 'w') as result_writer:
|
|
traceback.print_exc(file=result_writer)
|
|
# Prevent blocking in parent process by opening our end of all
|
|
# FIFOs.
|
|
for path, write in zip(fifo_paths, [True, False, False]):
|
|
_open_fifo(path, write)
|
|
os._exit(1)
|
|
with os.fdopen(result_w, 'w') as result_writer:
|
|
result_writer.write(f'{p.exitcode}\n')
|
|
|
|
os._exit(0)
|
|
else:
|
|
os.close(result_w)
|
|
fcntl.ioctl(tty, termios.TIOCSWINSZ,
|
|
struct.pack('HHHH', lines, columns, 0, 0))
|
|
env = IsolationEnvironment(env_pid, tty, *fifo_paths)
|
|
time.sleep(0.01)
|
|
try:
|
|
yield env
|
|
finally:
|
|
def get_return_code() -> int:
|
|
with os.fdopen(result_r) as result_reader:
|
|
result = result_reader.readline().rstrip()
|
|
try:
|
|
return int(result)
|
|
except ValueError:
|
|
pass
|
|
trace = result_reader.read()
|
|
raise TestProcessNotComplete('\n'.join([result,
|
|
trace]))
|
|
|
|
env.close(get_return_code)
|
|
|
|
|
|
def _run_test(test_function: typing.Callable[[], int],
|
|
pts: str,
|
|
stdin_fifo: Optional[str],
|
|
stdout_fifo: Optional[str],
|
|
stderr_fifo: Optional[str]) -> typing.NoReturn:
|
|
os.environ['TERM'] = 'xterm-256color'
|
|
|
|
tty_fd = os.open(pts, os.O_RDWR)
|
|
if stdin_fifo is not None:
|
|
old_stdin, sys.stdin = sys.stdin, open(stdin_fifo)
|
|
old_stdin.close()
|
|
else:
|
|
os.dup2(tty_fd, pty.STDIN_FILENO)
|
|
if stdout_fifo is not None:
|
|
old_stdout, sys.stdout = sys.stdout, open(stdout_fifo, 'w')
|
|
old_stdout.close()
|
|
else:
|
|
os.dup2(tty_fd, pty.STDOUT_FILENO)
|
|
if stderr_fifo is not None:
|
|
old_stderr, sys.stderr = sys.stderr, open(stderr_fifo, 'w',
|
|
buffering=1)
|
|
old_stderr.close()
|
|
else:
|
|
os.dup2(tty_fd, pty.STDERR_FILENO)
|
|
os.close(tty_fd)
|
|
|
|
result = test_function()
|
|
sys.exit(result or 0)
|