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())
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)