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)