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