• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2015 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4"""A module to abstract the shell execution environment on DUT."""
5
6import subprocess
7
8import time
9
10from autotest_lib.client.common_lib import error
11from autotest_lib.client.common_lib import utils
12
13
14class UnsupportedSuccessToken(Exception):
15    """Unsupported character found."""
16    pass
17
18
19class LocalShell(object):
20    """An object to wrap the local shell environment."""
21
22    def __init__(self, os_if):
23        """Initialize the LocalShell object."""
24        self._os_if = os_if
25
26    def _run_command(self, cmd, block=True):
27        """Helper function of run_command() methods.
28
29        Return the subprocess.Popen() instance to provide access to console
30        output in case command succeeded.  If block=False, will not wait for
31        process to return before returning.
32        """
33        self._os_if.log('Executing: %s' % cmd)
34        process = subprocess.Popen(
35                cmd,
36                shell=True,
37                stdout=subprocess.PIPE,
38                stderr=subprocess.PIPE)
39        if block:
40            process.wait()
41        return process
42
43    def run_command(self, cmd, block=True):
44        """Run a shell command.
45
46        In case of the command returning an error print its stdout and stderr
47        outputs on the console and dump them into the log. Otherwise suppress
48        all output.
49
50        @param block: if True (default), wait for command to finish
51        @raise error.CmdError: if block is True and command fails (rc!=0)
52        """
53        start_time = time.time()
54        process = self._run_command(cmd, block)
55        if block and process.returncode:
56            # Grab output only if an error occurred
57            returncode = process.returncode
58            stdout = process.stdout.read()
59            stderr = process.stderr.read()
60            duration = time.time() - start_time
61            result = utils.CmdResult(cmd, stdout, stderr, returncode, duration)
62            self._os_if.log('Command failed.\n%s' % result)
63            raise error.CmdError(cmd, result)
64
65    def run_command_get_result(self, cmd, ignore_status=False):
66        """Run a shell command, and get the result (output and returncode).
67
68        @param ignore_status: if True, do not raise CmdError, even if rc != 0.
69        @raise error.CmdError: if command fails (rc!=0) and not ignore_result
70        @return the result of the command
71        @rtype: utils.CmdResult
72        """
73        start_time = time.time()
74
75        process = self._run_command(cmd, block=True)
76
77        returncode = process.returncode
78        stdout = process.stdout.read()
79        stderr = process.stderr.read()
80        duration = time.time() - start_time
81        result = utils.CmdResult(cmd, stdout, stderr, returncode, duration)
82
83        if returncode and not ignore_status:
84            self._os_if.log('Command failed:\n%s' % result)
85            raise error.CmdError(cmd, result)
86
87        self._os_if.log('Command result:\n%s' % result)
88        return result
89
90    def run_command_check_output(self, cmd, success_token):
91        """Run a command and check whether standard output contains some string.
92
93        The sucess token is assumed to not contain newlines.
94
95        @param cmd: A string of the command to make a blocking call with.
96        @param success_token: A string to search the standard output of the
97                command for.
98
99        @returns a Boolean indicating whthere the success_token was in the
100                stdout of the cmd.
101
102        @raises UnsupportedSuccessToken if a newline is found in the
103                success_token.
104        """
105        # The run_command_get_outuput method strips newlines from stdout.
106        if '\n' in success_token:
107            raise UnsupportedSuccessToken()
108        cmd_stdout = ''.join(self.run_command_get_output(cmd))
109        self._os_if.log('Checking for %s in %s' % (success_token, cmd_stdout))
110        return success_token in cmd_stdout
111
112    def run_command_get_status(self, cmd):
113        """Run a shell command and return its return code.
114
115        The return code of the command is returned, in case of any error.
116        """
117        process = self._run_command(cmd)
118        return process.returncode
119
120    def run_command_get_output(self, cmd, include_stderr=False):
121        """Run shell command and return stdout (and possibly stderr) to the caller.
122
123        The output is returned as a list of strings stripped of the newline
124        characters.
125        """
126        process = self._run_command(cmd)
127        text = [x.rstrip() for x in process.stdout.readlines()]
128        if include_stderr:
129            text.extend([x.rstrip() for x in process.stderr.readlines()])
130        return text
131
132    def read_file(self, path):
133        """Read the content of the file."""
134        with open(path) as f:
135            return f.read()
136
137    def write_file(self, path, data):
138        """Write the data to the file."""
139        with open(path, 'w') as f:
140            f.write(data)
141
142    def append_file(self, path, data):
143        """Append the data to the file."""
144        with open(path, 'a') as f:
145            f.write(data)
146