def test_background_job_err(self): """Test reading standard err. Launches a BackgroundJob that writes to standard error to see if it gets captured. """ job = background_job.BackgroundJob('python', stdin=PY_ERR_CODE) self.assertTrue(job.result.stderr.startswith('TEST'))
def test_background_job(self): """Test running a BackgroundJob Runs a simple BackgroundJob to and checks its standard output to see if it ran. """ job = background_job.BackgroundJob('echo TEST') self.assertTrue(job.result.stdout.startswith('TEST'))
def setup_master_ssh(self, timeout_seconds=5): """Sets up the master ssh connection. Sets up the inital master ssh connection if it has not already been started. Args: timeout_seconds: The time to wait for the master ssh connection to be made. Raises: SshTimeoutError: If the master ssh connection takes to long to start then a timeout error will be thrown. """ with self._lock: if self._background_job is not None: socket_path = self.socket_path if (not os.path.exists(socket_path) or self._background_jobsp.poll() is not None): logging.info('Master ssh connection to %s is down.', self.connection.construct_host_name()) self._cleanup_master_ssh() if self._background_job is None: # Create a shared socket in a temp location. self._master_ssh_tempdir = tempfile.mkdtemp( prefix='ssh-master') # Setup flags and options for running the master ssh # -N: Do not execute a remote command. # ControlMaster: Spawn a master connection. # ControlPath: The master connection socket path. extra_flags = {'-N': None} extra_options = { 'ControlMaster': True, 'ControlPath': self.socket_path, 'BatchMode': True } # Construct the command and start it. master_cmd = self._formatter.format_ssh_local_command( self._settings, extra_flags, extra_options) logging.info('Starting master ssh connection %s', master_cmd) self._background_job = background_job.BackgroundJob( master_cmd, no_pipes=True) end_time = time.time() + timeout_seconds while time.time() < end_time: if os.path.exists(self.socket_path): break time.sleep(.2) else: raise error.SshTimeoutError( 'Master ssh connection timed out.')
def test_background_job_pipe(self): """Test piping on a BackgroundJob. Tests that the standard output of a job can be piped to another stream. """ mem_buffer = io.StringIO() job = background_job.BackgroundJob('echo TEST', stdout_tee=mem_buffer) job.wait() self.assertTrue(mem_buffer.getvalue().startswith('TEST'))
def test_background_job_in(self): """Test if a background job can send through standard input. Sends a line through standard input to see if the BackgroundJob can pick it up. """ job = background_job.BackgroundJob('grep "TEST"', allow_send=True) job.sendline('TEST') self.assertTrue(job.result.stdout.startswith('TEST'))
def test_background_job_pipe_err(self): """Test error piping on a BackgroundJob. Tests that the standard output of a job can be piped to another stream. """ mem_buffer = io.StringIO() job = background_job.BackgroundJob("python", stdin=PY_ERR_CODE, stderr_tee=mem_buffer) job.wait() self.assertTrue(mem_buffer.getvalue().startswith('TEST'))
def test_background_job_instream(self): """Test if a background job can pipe its stdin. Sends standard input to the BackgroundJob through a different stream. """ stream = io.BytesIO() job = background_job.BackgroundJob('grep "TEST"', stdin=stream) # In a real situation a pipe would probably be used, however writing # and then seeking is simpiler for the sake of testing. stream.write('TEST'.encode()) stream.seek(0) self.assertTrue(job.result.stdout.startswith('TEST'))
def run(self, command, timeout_seconds=3600, env={}, stdout=None, stderr=None, stdin=None, master_connection_timeout=5): """Run a remote command over ssh. Runs a remote command over ssh. Args: command: The command to execute over ssh. Can be either a string or a list. timeout_seconds: How long to wait on the command before timing out. env: A dictonary of enviroment variables to setup on the remote host. stdout: A stream to send stdout to. stderr: A stream to send stderr to. stdin: A string that contains the contents of stdin. A string may also be used, however its contents are not gurenteed to be sent to the ssh process. master_connection_timeout: The amount of time to wait for the master ssh connection to come up. Returns: The results of the ssh background job. Raises: CmdTimeoutError: When the remote command took to long to execute. SshTimeoutError: When the connection took to long to established. SshPermissionDeniedError: When permission is not allowed on the remote host. """ try: self.setup_master_ssh(master_connection_timeout) except error.SshError: logging.warning('Failed to create master ssh connection, using ' 'normal ssh connection.') extra_options = {'BatchMode': True} terminal_command = self._formatter.format_command( command, env, self._settings, extra_options=extra_options) dns_retry_count = 2 while True: job = background_job.BackgroundJob(terminal_command, stdout_tee=stdout, stderr_tee=stderr, verbose=False, stdin=stdin) job.wait(timeout_seconds) result = job.result error_string = result.stderr dns_retry_count -= 1 if (result and result.exit_status == 255 and re.search( r'^ssh: .*: Name or service not known', error_string)): if dns_retry_count: logging.debug('Retrying because of DNS failure') continue logging.debug('Retry failed.') elif not dns_retry_count: logging.debug('Retry succeeded.') break # The error messages will show up in band (indistinguishable # from stuff sent through the SSH connection), so we have the # remote computer echo the message "Connected." before running # any command. Since the following 2 errors have to do with # connecting, it's safe to do these checks. # This may not be true in acts? if result.exit_status == 255: if re.search( r'^ssh: connect to host .* port .*: ' r'Connection timed out\r$', error_string): raise error.SshTimeoutError('ssh timed out', result) if 'Permission denied' in error_string: msg = 'ssh permission denied' raise error.SshPermissionDeniedError(msg, result) if re.search( r'ssh: Could not resolve hostname .*: ' r'Name or service not known', error_string): raise error.SshUnknownHost('unknown host', result) return result
def test_background_job_env(self): job = background_job.BackgroundJob('printenv', env={'MYTESTVAR': '20'}) self.assertNotEqual(job.result.stdout.find('MYTESTVAR=20'), -1)
def test_background_job_timeout(self): with self.assertRaises(background_job.CmdTimeoutError): job = background_job.BackgroundJob('sleep 5') job.wait(timeout=0.1)