class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" def __init__(self, sock, in_fd, out_fd, err_fd): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock) if in_fd else None self._stdout = out_fd self._stderr = err_fd self.remote_pid = None def _maybe_start_input_reader(self): if self._input_reader: self._input_reader.start() def _maybe_stop_input_reader(self): if self._input_reader: self._input_reader.stop() 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._stdout.write(payload) self._stdout.flush() elif chunk_type == ChunkType.STDERR: self._stderr.write(payload) self._stderr.flush() elif chunk_type == ChunkType.EXIT: self._stdout.flush() self._stderr.flush() return int(payload) elif chunk_type == ChunkType.PID: self.remote_pid = int(payload) elif chunk_type == ChunkType.START_READING_INPUT: self._maybe_start_input_reader() 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 NailgunStreamReader is always stopped. self._maybe_stop_input_reader() 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()
class TestNailgunStreamReader(unittest.TestCase): def setUp(self): self.in_fd = FakeFile() self.mock_socket = mock.Mock() self.reader = NailgunStreamReader(in_fd=self.in_fd, sock=self.mock_socket) def test_stop(self): self.assertFalse(self.reader.is_stopped) self.reader.stop() self.assertTrue(self.reader.is_stopped) self.reader.run() def test_startable(self): self.assertTrue(inspect.ismethod(self.reader.start)) @mock.patch('select.select') def test_run_stop_on_error(self, mock_select): mock_select.return_value = ([], [], [self.in_fd]) self.reader.run() self.assertTrue(self.reader.is_stopped) 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 NailgunStreamReader.running() and .run() simultaneously. with self.reader.running(): while not self.reader.is_stopped: time.sleep(0.01) self.assertTrue(self.reader.is_stopped) mock_read.assert_called_with(-1, io.DEFAULT_BUFFER_SIZE) self.assertEquals(mock_read.call_count, 2) self.mock_socket.shutdown.assert_called_once_with(socket.SHUT_WR) mock_writer.assert_has_calls([ mock.call(mock.ANY, ChunkType.STDIN, b'A' * 300), mock.call(mock.ANY, ChunkType.STDIN_EOF) ])
class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" def __init__(self, sock, in_fd, out_fd, err_fd): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock) if in_fd else None self._stdout = out_fd self._stderr = err_fd self.remote_pid = None def _maybe_start_input_reader(self): if self._input_reader: self._input_reader.start() def _maybe_stop_input_reader(self): if self._input_reader: self._input_reader.stop() def _process_session(self): """Process the outputs of the nailgun session.""" try: for chunk_type, payload in self.iter_chunks(self._sock): if chunk_type == ChunkType.STDOUT: self._stdout.write(payload) self._stdout.flush() elif chunk_type == ChunkType.STDERR: self._stderr.write(payload) self._stderr.flush() elif chunk_type == ChunkType.EXIT: self._stdout.flush() self._stderr.flush() return int(payload) elif chunk_type == ChunkType.PID: self.remote_pid = int(payload) elif chunk_type == ChunkType.START_READING_INPUT: self._maybe_start_input_reader() 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 NailgunStreamReader is always stopped. self._maybe_stop_input_reader() 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()
class TestNailgunStreamReader(unittest.TestCase): def setUp(self): self.in_fd = FakeFile() self.mock_socket = mock.Mock() self.reader = NailgunStreamReader(in_fd=self.in_fd, sock=self.mock_socket) def test_stop(self): self.assertFalse(self.reader.is_stopped) self.reader.stop() self.assertTrue(self.reader.is_stopped) self.reader.run() def test_startable(self): self.assertTrue(inspect.ismethod(self.reader.start)) @mock.patch("select.select") def test_run_stop_on_error(self, mock_select): mock_select.return_value = ([], [], [self.in_fd]) self.reader.run() self.assertTrue(self.reader.is_stopped) 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 NailgunStreamReader.running() and .run() simultaneously. with self.reader.running(): while not self.reader.is_stopped: time.sleep(0.01) self.assertTrue(self.reader.is_stopped) mock_read.assert_called_with(-1, io.DEFAULT_BUFFER_SIZE) self.assertEquals(mock_read.call_count, 2) self.mock_socket.shutdown.assert_called_once_with(socket.SHUT_WR) mock_writer.assert_has_calls( [mock.call(mock.ANY, ChunkType.STDIN, b"A" * 300), mock.call(mock.ANY, ChunkType.STDIN_EOF)] )
class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" BUF_SIZE = 8192 def __init__(self, sock, in_fd, out_fd, err_fd): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock, self.BUF_SIZE) if in_fd else None self._stdout = out_fd self._stderr = err_fd @contextmanager def _maybe_input_reader_running(self): if self._input_reader: self._input_reader.start() yield if self._input_reader: self._input_reader.stop() def _process_session(self): """Process the outputs of the nailgun session.""" for chunk_type, payload in self.iter_chunks(self._sock): if chunk_type == ChunkType.STDOUT: self._stdout.write(payload) self._stdout.flush() elif chunk_type == ChunkType.STDERR: self._stderr.write(payload) self._stderr.flush() elif chunk_type == ChunkType.EXIT: self._stdout.flush() self._stderr.flush() return int(payload) else: raise self.ProtocolError('Received unexpected chunk {} -> {}'.format(chunk_type, payload)) def execute(self, working_dir, main_class, *arguments, **environment): # Send the nailgun request. self.send_request(self._sock, working_dir, main_class, *arguments, **environment) # Launch the NailgunStreamReader if applicable and process the remainder of the nailgun session. with self._maybe_input_reader_running(): return self._process_session()
class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" def __init__(self, sock, in_fd, out_fd, err_fd, exit_on_broken_pipe=False): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock) if in_fd else None self._stdout = out_fd self._stderr = err_fd self._exit_on_broken_pipe = exit_on_broken_pipe self.remote_pid = None def _maybe_start_input_reader(self): if self._input_reader: self._input_reader.start() def _maybe_stop_input_reader(self): if self._input_reader: self._input_reader.stop() 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_reader() 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 NailgunStreamReader is always stopped. self._maybe_stop_input_reader() 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()