#!/usr/bin/env python
"""
Tests the buffering behavior of stdio streams in `mercurial.utils.procutil`.
"""
import contextlib
import errno
import os
import pickle
import signal
import subprocess
import sys
import tempfile
import unittest
from mercurial import pycompat, util
TEST_BUFFERING_CHILD_SCRIPT = r'''
import os
from mercurial import dispatch
from mercurial.utils import procutil
dispatch.initstdio()
procutil.{stream}.write(b'aaa')
os.write(procutil.{stream}.fileno(), b'[written aaa]')
procutil.{stream}.write(b'bbb\n')
os.write(procutil.{stream}.fileno(), b'[written bbb\\n]')
'''
UNBUFFERED = b'aaa[written aaa]bbb\n[written bbb\\n]'
LINE_BUFFERED = b'[written aaa]aaabbb\n[written bbb\\n]'
FULLY_BUFFERED = b'[written aaa][written bbb\\n]aaabbb\n'
TEST_LARGE_WRITE_CHILD_SCRIPT = r'''
import os
import signal
import sys
from mercurial import dispatch
from mercurial.utils import procutil
signal.signal(signal.SIGINT, lambda *x: None)
dispatch.initstdio()
write_result = procutil.{stream}.write(b'x' * 1048576)
with os.fdopen(
os.open({write_result_fn!r}, os.O_WRONLY | getattr(os, 'O_TEMPORARY', 0)),
'w',
) as write_result_f:
write_result_f.write(str(write_result))
'''
TEST_BROKEN_PIPE_CHILD_SCRIPT = r'''
import os
import pickle
from mercurial import dispatch
from mercurial.utils import procutil
dispatch.initstdio()
procutil.stdin.read(1) # wait until parent process closed pipe
try:
procutil.{stream}.write(b'test')
procutil.{stream}.flush()
except EnvironmentError as e:
with os.fdopen(
os.open(
{err_fn!r},
os.O_WRONLY
| getattr(os, 'O_BINARY', 0)
| getattr(os, 'O_TEMPORARY', 0),
),
'wb',
) as err_f:
pickle.dump(e, err_f)
# Exit early to suppress further broken pipe errors at interpreter shutdown.
os._exit(0)
'''
@contextlib.contextmanager
def _closing(fds):
try:
yield
finally:
for fd in fds:
try:
os.close(fd)
except EnvironmentError:
pass
# In the following, we set the FDs non-inheritable mainly to make it possible
# for tests to close the receiving end of the pipe / PTYs.
@contextlib.contextmanager
def _devnull():
devnull = os.open(os.devnull, os.O_WRONLY)
# We don't have a receiving end, so it's not worth the effort on Python 2
# on Windows to make the FD non-inheritable.
with _closing([devnull]):
yield (None, devnull)
@contextlib.contextmanager
def _pipes():
rwpair = os.pipe()
with _closing(rwpair):
yield rwpair
@contextlib.contextmanager
def _ptys():
if pycompat.iswindows:
raise unittest.SkipTest("PTYs are not supported on Windows")
import pty
import tty
rwpair = pty.openpty()
with _closing(rwpair):
tty.setraw(rwpair[0])
yield rwpair
def _readall(fd, buffer_size, initial_buf=None):
buf = initial_buf or []
while True:
try:
s = os.read(fd, buffer_size)
except OSError as e:
if e.errno == errno.EIO:
# If the child-facing PTY got closed, reading from the
# parent-facing PTY raises EIO.
break
raise
if not s:
break
buf.append(s)
return b''.join(buf)
class TestStdio(unittest.TestCase):
def _test(
self,
child_script,
stream,
rwpair_generator,
check_output,
python_args=[],
post_child_check=None,
stdin_generator=None,
):
assert stream in ('stdout', 'stderr')
if stdin_generator is None:
stdin_generator = open(os.devnull, 'rb')
with rwpair_generator() as (
stream_receiver,
child_stream,
), stdin_generator as child_stdin:
proc = subprocess.Popen(
[sys.executable] + python_args + ['-c', child_script],
stdin=child_stdin,
stdout=child_stream if stream == 'stdout' else None,
stderr=child_stream if stream == 'stderr' else None,
)
try:
os.close(child_stream)
if stream_receiver is not None:
check_output(stream_receiver, proc)
except: # re-raises
proc.terminate()
raise
finally:
retcode = proc.wait()
self.assertEqual(retcode, 0)
if post_child_check is not None:
post_child_check()
def _test_buffering(
self, stream, rwpair_generator, expected_output, python_args=[]
):
def check_output(stream_receiver, proc):
self.assertEqual(_readall(stream_receiver, 1024), expected_output)
self._test(
TEST_BUFFERING_CHILD_SCRIPT.format(stream=stream),
stream,
rwpair_generator,
check_output,
python_args,
)
def test_buffering_stdout_devnull(self):
self._test_buffering('stdout', _devnull, None)
def test_buffering_stdout_pipes(self):
self._test_buffering('stdout', _pipes, FULLY_BUFFERED)
def test_buffering_stdout_ptys(self):
self._test_buffering('stdout', _ptys, LINE_BUFFERED)
def test_buffering_stdout_devnull_unbuffered(self):
self._test_buffering('stdout', _devnull, None, python_args=['-u'])
def test_buffering_stdout_pipes_unbuffered(self):
self._test_buffering('stdout', _pipes, UNBUFFERED, python_args=['-u'])
def test_buffering_stdout_ptys_unbuffered(self):
self._test_buffering('stdout', _ptys, UNBUFFERED, python_args=['-u'])
def _test_large_write(self, stream, rwpair_generator, python_args=[]):
def check_output(stream_receiver, proc):
if not pycompat.iswindows:
# On Unix, we can provoke a partial write() by interrupting it
# by a signal handler as soon as a bit of data was written.
# We test that write() is called until all data is written.
buf = [os.read(stream_receiver, 1)]
proc.send_signal(signal.SIGINT)
else:
# On Windows, there doesn't seem to be a way to cause partial
# writes.
buf = []
self.assertEqual(
_readall(stream_receiver, 131072, buf), b'x' * 1048576
)
def post_child_check():
self.assertEqual(write_result_f.read(), '1048576')
with tempfile.NamedTemporaryFile('r') as write_result_f:
self._test(
TEST_LARGE_WRITE_CHILD_SCRIPT.format(
stream=stream, write_result_fn=write_result_f.name
),
stream,
rwpair_generator,
check_output,
python_args,
post_child_check=post_child_check,
)
def test_large_write_stdout_devnull(self):
self._test_large_write('stdout', _devnull)
def test_large_write_stdout_pipes(self):
self._test_large_write('stdout', _pipes)
def test_large_write_stdout_ptys(self):
self._test_large_write('stdout', _ptys)
def test_large_write_stdout_devnull_unbuffered(self):
self._test_large_write('stdout', _devnull, python_args=['-u'])
def test_large_write_stdout_pipes_unbuffered(self):
self._test_large_write('stdout', _pipes, python_args=['-u'])
def test_large_write_stdout_ptys_unbuffered(self):
self._test_large_write('stdout', _ptys, python_args=['-u'])
def test_large_write_stderr_devnull(self):
self._test_large_write('stderr', _devnull)
def test_large_write_stderr_pipes(self):
self._test_large_write('stderr', _pipes)
def test_large_write_stderr_ptys(self):
self._test_large_write('stderr', _ptys)
def test_large_write_stderr_devnull_unbuffered(self):
self._test_large_write('stderr', _devnull, python_args=['-u'])
def test_large_write_stderr_pipes_unbuffered(self):
self._test_large_write('stderr', _pipes, python_args=['-u'])
def test_large_write_stderr_ptys_unbuffered(self):
self._test_large_write('stderr', _ptys, python_args=['-u'])
def _test_broken_pipe(self, stream):
assert stream in ('stdout', 'stderr')
def check_output(stream_receiver, proc):
os.close(stream_receiver)
proc.stdin.write(b'x')
proc.stdin.close()
def post_child_check():
err = pickle.load(err_f)
self.assertEqual(err.errno, errno.EPIPE)
self.assertEqual(err.strerror, "Broken pipe")
with tempfile.NamedTemporaryFile('rb') as err_f:
self._test(
TEST_BROKEN_PIPE_CHILD_SCRIPT.format(
stream=stream, err_fn=err_f.name
),
stream,
_pipes,
check_output,
post_child_check=post_child_check,
stdin_generator=util.nullcontextmanager(subprocess.PIPE),
)
def test_broken_pipe_stdout(self):
self._test_broken_pipe('stdout')
def test_broken_pipe_stderr(self):
self._test_broken_pipe('stderr')
if __name__ == '__main__':
import silenttestrunner
silenttestrunner.main(__name__)