class TestNailgunStreamWriter(unittest.TestCase):
TEST_VALUE = "1729"
def setUp(self):
self.chunk_type = ChunkType.STDERR
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type)
@mock.patch.object(NailgunProtocol, "write_chunk")
def test_write(self, mock_writer):
self.writer.write(self.TEST_VALUE)
mock_writer.assert_called_once_with(self.mock_socket, self.chunk_type, self.TEST_VALUE)
@mock.patch.object(NailgunProtocol, "write_chunk")
def test_write_broken_pipe_unmasked(self, mock_writer):
mock_writer.side_effect = IOError(errno.EPIPE, os.strerror(errno.EPIPE))
with self.assertRaises(IOError):
self.writer.write(self.TEST_VALUE)
@mock.patch.object(NailgunProtocol, "write_chunk")
def test_write_broken_pipe_masked(self, mock_writer):
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type, mask_broken_pipe=True)
mock_writer.side_effect = IOError(errno.EPIPE, os.strerror(errno.EPIPE))
self.writer.write(self.TEST_VALUE)
def test_isatty(self):
self.assertTrue(self.writer.isatty())
def test_not_isatty(self):
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type, isatty=False)
self.assertFalse(self.writer.isatty())
def test_misc(self):
self.writer.flush()
def test_write_broken_pipe_masked(self, mock_writer):
self.writer = NailgunStreamWriter(self.mock_socket,
self.chunk_type,
mask_broken_pipe=True)
mock_writer.side_effect = IOError(errno.EPIPE,
os.strerror(errno.EPIPE))
self.writer.write(self.TEST_VALUE)
class TestNailgunStreamWriter(unittest.TestCase):
def setUp(self):
self.in_fd = -1
self.mock_socket = unittest.mock.Mock()
self.writer = NailgunStreamWriter(
(self.in_fd,),
self.mock_socket,
(ChunkType.STDIN,),
ChunkType.STDIN_EOF
)
def test_stop(self):
self.assertFalse(self.writer.is_stopped)
self.writer.stop()
self.assertTrue(self.writer.is_stopped)
self.writer.run()
def test_startable(self):
self.assertTrue(inspect.ismethod(self.writer.start))
@unittest.mock.patch('select.select')
def test_run_stop_on_error(self, mock_select):
mock_select.return_value = ([], [], [self.in_fd])
self.writer.run()
self.assertFalse(self.writer.is_alive())
self.assertEqual(mock_select.call_count, 1)
@unittest.mock.patch('os.read')
@unittest.mock.patch('select.select')
@unittest.mock.patch.object(NailgunProtocol, 'write_chunk')
def test_run_read_write(self, mock_writer, mock_select, mock_read):
mock_select.side_effect = [
([self.in_fd], [], []),
([self.in_fd], [], [])
]
mock_read.side_effect = [
b'A' * 300,
b'' # Simulate EOF.
]
# Exercise NailgunStreamWriter.running() and .run() simultaneously.
inc = 0
with self.writer.running():
while self.writer.is_alive():
time.sleep(0.01)
inc += 1
if inc >= 1000:
raise Exception('waited too long.')
self.assertFalse(self.writer.is_alive())
mock_read.assert_called_with(-1, io.DEFAULT_BUFFER_SIZE)
self.assertEqual(mock_read.call_count, 2)
mock_writer.assert_has_calls([
unittest.mock.call(unittest.mock.ANY, ChunkType.STDIN, b'A' * 300),
unittest.mock.call(unittest.mock.ANY, ChunkType.STDIN_EOF)
])
class TestNailgunStreamWriter(unittest.TestCase):
def setUp(self):
self.in_fd = -1
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(
(self.in_fd,),
self.mock_socket,
(ChunkType.STDIN,),
ChunkType.STDIN_EOF
)
def test_stop(self):
self.assertFalse(self.writer.is_stopped)
self.writer.stop()
self.assertTrue(self.writer.is_stopped)
self.writer.run()
def test_startable(self):
self.assertTrue(inspect.ismethod(self.writer.start))
@mock.patch('select.select')
def test_run_stop_on_error(self, mock_select):
mock_select.return_value = ([], [], [self.in_fd])
self.writer.run()
self.assertFalse(self.writer.is_alive())
self.assertEquals(mock_select.call_count, 1)
@mock.patch('os.read')
@mock.patch('select.select')
@mock.patch.object(NailgunProtocol, 'write_chunk')
def test_run_read_write(self, mock_writer, mock_select, mock_read):
mock_select.side_effect = [
([self.in_fd], [], []),
([self.in_fd], [], [])
]
mock_read.side_effect = [
b'A' * 300,
b'' # Simulate EOF.
]
# Exercise NailgunStreamWriter.running() and .run() simultaneously.
inc = 0
with self.writer.running():
while self.writer.is_alive():
time.sleep(0.01)
inc += 1
if inc >= 1000:
raise Exception('waited too long.')
self.assertFalse(self.writer.is_alive())
mock_read.assert_called_with(-1, io.DEFAULT_BUFFER_SIZE)
self.assertEquals(mock_read.call_count, 2)
mock_writer.assert_has_calls([
mock.call(mock.ANY, ChunkType.STDIN, b'A' * 300),
mock.call(mock.ANY, ChunkType.STDIN_EOF)
])
def setUp(self):
self.in_fd = -1
self.mock_socket = unittest.mock.Mock()
self.writer = NailgunStreamWriter(
(self.in_fd,),
self.mock_socket,
(ChunkType.STDIN,),
ChunkType.STDIN_EOF
)
def __init__(self, sock, in_fd, out_fd, err_fd, exit_on_broken_pipe=False):
self._sock = sock
if in_fd:
self._input_writer = NailgunStreamWriter(in_fd, self._sock,
ChunkType.STDIN,
ChunkType.STDIN_EOF)
else:
self._input_writer = None
self._stdout = out_fd
self._stderr = err_fd
self._exit_on_broken_pipe = exit_on_broken_pipe
self.remote_pid = None
def _nailgunned_stdio(self, sock):
"""Redirects stdio to the connected socket speaking the nailgun protocol."""
# Determine output tty capabilities from the environment.
stdin_isatty, stdout_isatty, stderr_isatty = NailgunProtocol.isatty_from_env(
self._env)
# Launch a thread to read stdin data from the socket (the only messages expected from the client
# for the remainder of the protocol), and threads to copy from stdout/stderr pipes onto the
# socket.
with NailgunStreamStdinReader.open(sock, isatty=stdin_isatty) as stdin,\
NailgunStreamWriter.open(sock, ChunkType.STDOUT, None, isatty=stdout_isatty) as stdout,\
NailgunStreamWriter.open(sock, ChunkType.STDERR, None, isatty=stderr_isatty) as stderr:
with stdio_as(stdout=stdout, stderr=stderr, stdin=stdin):
yield
def __init__(self, sock, in_file, out_file, err_file, exit_on_broken_pipe=False):
self._sock = sock
self._input_writer = None
if in_file:
self._input_writer = NailgunStreamWriter(
(in_file.fileno(),),
self._sock,
(ChunkType.STDIN,),
ChunkType.STDIN_EOF
)
self._stdout = out_file
self._stderr = err_file
self._exit_on_broken_pipe = exit_on_broken_pipe
self.remote_pid = None
def __init__(
self,
sock,
in_file,
out_file,
err_file,
exit_on_broken_pipe=False,
remote_pid_callback=None,
remote_pgrp_callback=None,
):
"""
:param bool exit_on_broken_pipe: whether or not to exit when `Broken Pipe` errors are
encountered
:param remote_pid_callback: Callback to run when a pid chunk is received from a remote client.
:param remote_pgrp_callback: Callback to run when a pgrp (process group) chunk is received from
a remote client.
"""
self._sock = sock
self._input_writer = (None if not in_file else NailgunStreamWriter(
(in_file.fileno(), ), self._sock,
(ChunkType.STDIN, ), ChunkType.STDIN_EOF))
self._stdout = out_file
self._stderr = err_file
self._exit_on_broken_pipe = exit_on_broken_pipe
self.remote_pid = None
self.remote_process_cmdline = None
self.remote_pgrp = None
self._remote_pid_callback = remote_pid_callback
self._remote_pgrp_callback = remote_pgrp_callback
# NB: These variables are set in a signal handler to implement graceful shutdown.
self._exit_timeout_start_time = None
self._exit_timeout = None
self._exit_reason = None
def _pipe_stdio(self, sock, stdin_isatty, stdout_isatty, stderr_isatty):
"""Handles stdio redirection in the case of pipes and/or mixed pipes and ttys."""
stdio_writers = ((ChunkType.STDOUT, stdout_isatty), (ChunkType.STDERR,
stderr_isatty))
types, ttys = zip(*(stdio_writers))
with NailgunStreamStdinReader.open(sock, stdin_isatty) as stdin_fd,\
NailgunStreamWriter.open_multi(sock, types, ttys) as ((stdout_fd, stderr_fd), writer),\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(
.001
) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def _nailgunned_stdio(self, sock):
"""Redirects stdio to the connected socket speaking the nailgun protocol."""
# Determine output tty capabilities from the environment.
stdin_isatty, stdout_isatty, stderr_isatty = NailgunProtocol.isatty_from_env(
self._env)
# Launch a thread to read stdin data from the socket (the only messages expected from the client
# for the remainder of the protocol), and threads to copy from stdout/stderr pipes onto the
# socket.
with NailgunStreamWriter.open_multi(
sock,
(ChunkType.STDOUT, ChunkType.STDERR),
None,
(stdout_isatty, stderr_isatty)
) as ((stdout_fd, stderr_fd), writer),\
NailgunStreamStdinReader.open(sock, stdin_isatty) as stdin_fd,\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(
.001
) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def __init__(self,
sock,
in_file,
out_file,
err_file,
exit_on_broken_pipe=False,
remote_pid_callback=None,
remote_pgrp_callback=None):
"""
:param bool exit_on_broken_pipe: whether or not to exit when `Broken Pipe` errors are
encountered
:param remote_pid_callback: Callback to run when a pid chunk is received from a remote client.
:param remote_pgrp_callback: Callback to run when a pgrp (process group) chunk is received from
a remote client.
"""
self._sock = sock
self._input_writer = None if not in_file else NailgunStreamWriter(
(in_file.fileno(), ), self._sock,
(ChunkType.STDIN, ), ChunkType.STDIN_EOF)
self._stdout = out_file
self._stderr = err_file
self._exit_on_broken_pipe = exit_on_broken_pipe
self.remote_pid = None
self.remote_pgrp = None
self._remote_pid_callback = remote_pid_callback
self._remote_pgrp_callback = remote_pgrp_callback
def _nailgunned_stdio(self, sock):
"""Redirects stdio to the connected socket speaking the nailgun protocol."""
# Determine output tty capabilities from the environment.
_, stdout_isatty, stderr_isatty = NailgunProtocol.isatty_from_env(self._env)
# TODO(kwlzn): Implement remote input reading and fix the non-fork()-safe sys.stdin reference
# in NailgunClient to enable support for interactive goals like `repl` etc.
# Construct StreamWriters for stdout, stderr.
streams = (
NailgunStreamWriter(sock, ChunkType.STDOUT, isatty=stdout_isatty),
NailgunStreamWriter(sock, ChunkType.STDERR, isatty=stderr_isatty)
)
# Launch the stdin StreamReader and redirect stdio.
with stdio_as(*streams):
yield
def setUp(self):
self.in_fd = -1
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(
(self.in_fd,),
self.mock_socket,
(ChunkType.STDIN,),
ChunkType.STDIN_EOF
)
def _nailgunned_stdio(self, sock):
"""Redirects stdio to the connected socket speaking the nailgun protocol."""
# Determine output tty capabilities from the environment.
_, stdout_isatty, stderr_isatty = NailgunProtocol.isatty_from_env(
self._env)
# Construct a StreamReader for stdin.
stdin_reader = NailgunStreamReader(sys.stdin, sock)
# Construct StreamWriters for stdout, stderr.
streams = (NailgunStreamWriter(sock,
ChunkType.STDOUT,
isatty=stdout_isatty),
NailgunStreamWriter(sock,
ChunkType.STDERR,
isatty=stderr_isatty), stdin_reader)
# Launch the stdin StreamReader and redirect stdio.
with stdin_reader.running(), stdio_as(*streams):
yield
class TestNailgunStreamWriter(unittest.TestCase):
TEST_VALUE = '1729'
def setUp(self):
self.chunk_type = ChunkType.STDERR
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type)
@mock.patch.object(NailgunProtocol, 'write_chunk')
def test_write(self, mock_writer):
self.writer.write(self.TEST_VALUE)
mock_writer.assert_called_once_with(self.mock_socket, self.chunk_type,
self.TEST_VALUE)
@mock.patch.object(NailgunProtocol, 'write_chunk')
def test_write_broken_pipe_unmasked(self, mock_writer):
mock_writer.side_effect = IOError(errno.EPIPE,
os.strerror(errno.EPIPE))
with self.assertRaises(IOError):
self.writer.write(self.TEST_VALUE)
@mock.patch.object(NailgunProtocol, 'write_chunk')
def test_write_broken_pipe_masked(self, mock_writer):
self.writer = NailgunStreamWriter(self.mock_socket,
self.chunk_type,
mask_broken_pipe=True)
mock_writer.side_effect = IOError(errno.EPIPE,
os.strerror(errno.EPIPE))
self.writer.write(self.TEST_VALUE)
def test_isatty(self):
self.assertTrue(self.writer.isatty())
def test_not_isatty(self):
self.writer = NailgunStreamWriter(self.mock_socket,
self.chunk_type,
isatty=False)
self.assertFalse(self.writer.isatty())
def test_misc(self):
self.writer.flush()
def _nailgunned_stdio(self, sock):
"""Redirects stdio to the connected socket speaking the nailgun protocol."""
# Determine output tty capabilities from the environment.
stdin_isatty, stdout_isatty, stderr_isatty = NailgunProtocol.isatty_from_env(
self._env)
# If all stdio is a tty, there's only one logical I/O device (the tty device). This happens to
# be addressable as a file in OSX and Linux, so we take advantage of that and directly open the
# character device for output redirection - eliminating the need to directly marshall any
# interactive stdio back/forth across the socket and permitting full, correct tty control with
# no middle-man.
if all((stdin_isatty, stdout_isatty, stderr_isatty)):
stdin_ttyname, stdout_ttyname, stderr_ttyname = NailgunProtocol.ttynames_from_env(
self._env)
assert stdin_ttyname == stdout_ttyname == stderr_ttyname, (
'expected all stdio ttys to be the same, but instead got: {}\n'
'please file a bug at http://github.com/pantsbuild/pants'.
format([stdin_ttyname, stdout_ttyname, stderr_ttyname]))
with open(stdin_ttyname, 'rb+wb', 0) as tty:
tty_fileno = tty.fileno()
with stdio_as(stdin_fd=tty_fileno,
stdout_fd=tty_fileno,
stderr_fd=tty_fileno):
def finalizer():
termios.tcdrain(tty_fileno)
yield finalizer
else:
stdio_writers = ((ChunkType.STDOUT, stdout_isatty),
(ChunkType.STDERR, stderr_isatty))
types, ttys = zip(*(stdio_writers))
with NailgunStreamStdinReader.open(sock, stdin_isatty) as stdin_fd,\
NailgunStreamWriter.open_multi(sock, types, ttys) as ((stdout_fd, stderr_fd), writer),\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(
.001
) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def __init__(self, sock, in_file, out_file, err_file, exit_on_broken_pipe=False):
"""
:param bool exit_on_broken_pipe: whether or not to exit when `Broken Pipe` errors are
encountered
"""
self._sock = sock
self._input_writer = (
None
if not in_file
else NailgunStreamWriter(
(in_file.fileno(),), self._sock, (ChunkType.STDIN,), ChunkType.STDIN_EOF
)
)
self._stdout = out_file
self._stderr = err_file
self._exit_on_broken_pipe = exit_on_broken_pipe
# NB: These variables are set in a signal handler to implement graceful shutdown.
self._exit_timeout_start_time = None
self._exit_timeout = None
self._exit_reason = None
def _pipe_stdio(cls, sock, stdin_isatty, stdout_isatty, stderr_isatty,
handle_stdin):
"""Handles stdio redirection in the case of pipes and/or mixed pipes and ttys."""
stdio_writers = ((ChunkType.STDOUT, stdout_isatty), (ChunkType.STDERR,
stderr_isatty))
types, ttys = zip(*(stdio_writers))
@contextmanager
def maybe_handle_stdin(want):
if want:
# TODO: Launching this thread pre-fork to handle @rule input currently results
# in an unhandled SIGILL in `src/python/pants/engine/scheduler.py, line 313 in pre_fork`.
# More work to be done here in https://github.com/pantsbuild/pants/issues/6005
with NailgunStreamStdinReader.open(sock, stdin_isatty) as fd:
yield fd
else:
with open('/dev/null', 'rb') as fh:
yield fh.fileno()
with maybe_handle_stdin(handle_stdin) as stdin_fd,\
NailgunStreamWriter.open_multi(sock, types, ttys) as ((stdout_fd, stderr_fd), writer),\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(
.001
) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def _pipe_stdio(cls, maybe_shutdown_socket, stdin_isatty, stdout_isatty,
stderr_isatty, handle_stdin):
"""Handles stdio redirection in the case of pipes and/or mixed pipes and ttys."""
stdio_writers = ((ChunkType.STDOUT, stdout_isatty), (ChunkType.STDERR,
stderr_isatty))
types, ttys = zip(*(stdio_writers))
@contextmanager
def maybe_handle_stdin(want):
if want:
with NailgunStreamStdinReader.open(maybe_shutdown_socket,
stdin_isatty) as fd:
yield fd
else:
with open('/dev/null', 'rb') as fh:
yield fh.fileno()
# TODO https://github.com/pantsbuild/pants/issues/7653
with maybe_handle_stdin(handle_stdin) as stdin_fd,\
NailgunStreamWriter.open_multi(maybe_shutdown_socket.socket, types, ttys) as ((stdout_fd, stderr_fd), writer),\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(
.001
) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def _pipe_stdio(cls, sock, stdin_isatty, stdout_isatty, stderr_isatty, handle_stdin):
"""Handles stdio redirection in the case of pipes and/or mixed pipes and ttys."""
stdio_writers = (
(ChunkType.STDOUT, stdout_isatty),
(ChunkType.STDERR, stderr_isatty)
)
types, ttys = zip(*(stdio_writers))
@contextmanager
def maybe_handle_stdin(want):
if want:
# TODO: Launching this thread pre-fork to handle @rule input currently results
# in an unhandled SIGILL in `src/python/pants/engine/scheduler.py, line 313 in pre_fork`.
# More work to be done here in https://github.com/pantsbuild/pants/issues/6005
with NailgunStreamStdinReader.open(sock, stdin_isatty) as fd:
yield fd
else:
with open('/dev/null', 'rb') as fh:
yield fh.fileno()
with maybe_handle_stdin(handle_stdin) as stdin_fd,\
NailgunStreamWriter.open_multi(sock, types, ttys) as ((stdout_fd, stderr_fd), writer),\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(.001) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def _pipe_stdio(self, sock, stdin_isatty, stdout_isatty, stderr_isatty):
"""Handles stdio redirection in the case of pipes and/or mixed pipes and ttys."""
stdio_writers = (
(ChunkType.STDOUT, stdout_isatty),
(ChunkType.STDERR, stderr_isatty)
)
types, ttys = zip(*(stdio_writers))
with NailgunStreamStdinReader.open(sock, stdin_isatty) as stdin_fd,\
NailgunStreamWriter.open_multi(sock, types, ttys) as ((stdout_fd, stderr_fd), writer),\
stdio_as(stdout_fd=stdout_fd, stderr_fd=stderr_fd, stdin_fd=stdin_fd):
# N.B. This will be passed to and called by the `DaemonExiter` prior to sending an
# exit chunk, to avoid any socket shutdown vs write races.
stdout, stderr = sys.stdout, sys.stderr
def finalizer():
try:
stdout.flush()
stderr.flush()
finally:
time.sleep(.001) # HACK: Sleep 1ms in the main thread to free the GIL.
writer.stop()
writer.join()
stdout.close()
stderr.close()
yield finalizer
def setUp(self):
self.chunk_type = ChunkType.STDERR
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type)
def test_not_isatty(self):
self.writer = NailgunStreamWriter(self.mock_socket,
self.chunk_type,
isatty=False)
self.assertFalse(self.writer.isatty())
def test_not_isatty(self):
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type, isatty=False)
self.assertFalse(self.writer.isatty())
def test_write_broken_pipe_masked(self, mock_writer):
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type, mask_broken_pipe=True)
mock_writer.side_effect = IOError(errno.EPIPE, os.strerror(errno.EPIPE))
self.writer.write(self.TEST_VALUE)
class TestNailgunStreamWriter(unittest.TestCase):
def setUp(self):
self.in_fd = -1
self.mock_socket = unittest.mock.Mock()
self.writer = NailgunStreamWriter((self.in_fd, ), self.mock_socket,
(ChunkType.STDIN, ),
ChunkType.STDIN_EOF)
def test_stop(self):
self.assertFalse(self.writer.is_stopped)
self.writer.stop()
self.assertTrue(self.writer.is_stopped)
self.writer.run()
def test_startable(self):
self.assertTrue(inspect.ismethod(self.writer.start))
@unittest.mock.patch("select.select")
def test_run_stop_on_error(self, mock_select):
mock_select.return_value = ([], [], [self.in_fd])
self.writer.run()
self.assertFalse(self.writer.is_alive())
self.assertEqual(mock_select.call_count, 1)
@unittest.mock.patch("os.read")
@unittest.mock.patch("select.select")
@unittest.mock.patch.object(NailgunProtocol, "write_chunk")
def test_run_read_write(self, mock_writer, mock_select, mock_read):
mock_select.side_effect = [([self.in_fd], [], []),
([self.in_fd], [], [])]
mock_read.side_effect = [b"A" * 300, b""] # Simulate EOF.
# Exercise NailgunStreamWriter.running() and .run() simultaneously.
inc = 0
with self.writer.running():
while self.writer.is_alive():
time.sleep(0.01)
inc += 1
if inc >= 1000:
raise Exception("waited too long.")
self.assertFalse(self.writer.is_alive())
mock_read.assert_called_with(-1, io.DEFAULT_BUFFER_SIZE)
self.assertEqual(mock_read.call_count, 2)
mock_writer.assert_has_calls([
unittest.mock.call(unittest.mock.ANY, ChunkType.STDIN, b"A" * 300),
unittest.mock.call(unittest.mock.ANY, ChunkType.STDIN_EOF),
])
@unittest.mock.patch("os.close")
@unittest.mock.patch("os.read")
@unittest.mock.patch("select.select")
def test_run_exits_for_closed_and_errored_socket(self, mock_select,
mock_read, mock_close):
# When stdin is closed, select can indicate that it is both empty and errored.
mock_select.return_value = ([self.in_fd], [self.in_fd], [])
mock_read.return_value = b"" # EOF.
self.writer.run()
assert self.writer.is_alive() is False
assert mock_select.call_count == 1
class NailgunClientSession(NailgunProtocol):
"""Handles a single nailgun client session."""
def __init__(self, sock, in_file, out_file, err_file, exit_on_broken_pipe=False):
self._sock = sock
self._input_writer = None
if in_file:
self._input_writer = NailgunStreamWriter(
(in_file.fileno(),),
self._sock,
(ChunkType.STDIN,),
ChunkType.STDIN_EOF
)
self._stdout = out_file
self._stderr = err_file
self._exit_on_broken_pipe = exit_on_broken_pipe
self.remote_pid = None
def _maybe_start_input_writer(self):
if self._input_writer:
self._input_writer.start()
def _maybe_stop_input_writer(self):
if self._input_writer and self._input_writer.is_alive():
self._input_writer.stop()
self._input_writer.join()
def _write_flush(self, fd, payload=None):
"""Write a payload to a given fd (if provided) and flush the fd."""
try:
if payload:
fd.write(payload)
fd.flush()
except (IOError, OSError) as e:
# If a `Broken Pipe` is encountered during a stdio fd write, we're headless - bail.
if e.errno == errno.EPIPE and self._exit_on_broken_pipe:
sys.exit()
# Otherwise, re-raise.
raise
def _process_session(self):
"""Process the outputs of the nailgun session."""
try:
for chunk_type, payload in self.iter_chunks(self._sock, return_bytes=True):
if chunk_type == ChunkType.STDOUT:
self._write_flush(self._stdout, payload)
elif chunk_type == ChunkType.STDERR:
self._write_flush(self._stderr, payload)
elif chunk_type == ChunkType.EXIT:
self._write_flush(self._stdout)
self._write_flush(self._stderr)
return int(payload)
elif chunk_type == ChunkType.PID:
self.remote_pid = int(payload)
elif chunk_type == ChunkType.START_READING_INPUT:
self._maybe_start_input_writer()
else:
raise self.ProtocolError('received unexpected chunk {} -> {}'.format(chunk_type, payload))
finally:
# Bad chunk types received from the server can throw NailgunProtocol.ProtocolError in
# NailgunProtocol.iter_chunks(). This ensures the NailgunStreamWriter is always stopped.
self._maybe_stop_input_writer()
def execute(self, working_dir, main_class, *arguments, **environment):
# Send the nailgun request.
self.send_request(self._sock, working_dir, main_class, *arguments, **environment)
# Process the remainder of the nailgun session.
return self._process_session()
def setUp(self):
self.in_file = FakeFile()
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(self.in_file, self.mock_socket,
ChunkType.STDIN, ChunkType.STDIN_EOF)
def setUp(self):
self.chunk_type = ChunkType.STDERR
self.mock_socket = mock.Mock()
self.writer = NailgunStreamWriter(self.mock_socket, self.chunk_type)
