xref: /aosp_15_r20/external/autotest/site_utils/lib/infra.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7) 
1# Copyright (c) 2014 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
5import contextlib
6import getpass
7import subprocess
8import os
9
10import common
11from autotest_lib.server.hosts import ssh_host
12from autotest_lib.client.common_lib import error
13from autotest_lib.client.common_lib import global_config
14from autotest_lib.client.common_lib import utils
15
16
17@contextlib.contextmanager
18def chdir(dirname=None):
19    """A context manager to help change directories.
20
21    Will chdir into the provided dirname for the lifetime of the context and
22    return to cwd thereafter.
23
24    @param dirname: The dirname to chdir into.
25    """
26    curdir = os.getcwd()
27    try:
28        if dirname is not None:
29            os.chdir(dirname)
30        yield
31    finally:
32        os.chdir(curdir)
33
34
35def local_runner(cmd, stream_output=False):
36    """
37    Runs a command on the local system as the current user.
38
39    @param cmd: The command to run.
40    @param stream_output: If True, streams the stdout of the process.
41
42    @returns: The output of cmd, will be stdout and stderr.
43    @raises CalledProcessError: If there was a non-0 return code.
44    """
45    print('Running command: %s' % cmd)
46    proc = subprocess.Popen(
47        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
48    if stream_output:
49        output = ''
50        for newline in iter(proc.stdout.readline, ''):
51            output += newline
52            print(newline.rstrip(os.linesep))
53    else:
54        output = proc.communicate()[0]
55
56    return_code = proc.wait()
57    if return_code !=0:
58        print("ERROR: '%s' failed with error:\n%s" % (cmd, output))
59        raise subprocess.CalledProcessError(return_code, cmd, output[:1024])
60    return output
61
62
63_host_objects = {}
64
65def host_object_runner(host, **kwargs):
66    """
67    Returns a function that returns the output of running a command via a host
68    object.
69
70    @param host: The host to run a command on.
71    @returns: A function that can invoke a command remotely.
72    """
73    try:
74        host_object = _host_objects[host]
75    except KeyError:
76        username = global_config.global_config.get_config_value(
77                'CROS', 'infrastructure_user')
78        host_object = ssh_host.SSHHost(host, user=username)
79        _host_objects[host] = host_object
80
81    def runner(cmd):
82        """
83        Runs a command via a host object on the enclosed host.  Translates
84        host.run errors to the subprocess equivalent to expose a common API.
85
86        @param cmd: The command to run.
87        @returns: The output of cmd.
88        @raises CalledProcessError: If there was a non-0 return code.
89        """
90        try:
91            return host_object.run(cmd).stdout
92        except error.AutotestHostRunError as e:
93            exit_status = e.result_obj.exit_status
94            command = e.result_obj.command
95            raise subprocess.CalledProcessError(exit_status, command)
96    return runner
97
98
99def googlesh_runner(host, **kwargs):
100    """
101    Returns a function that return the output of running a command via shelling
102    out to `googlesh`.
103
104    @param host: The host to run a command on
105    @returns: A function that can invoke a command remotely.
106    """
107    def runner(cmd):
108        """
109        Runs a command via googlesh on the enclosed host.
110
111        @param cmd: The command to run.
112        @returns: The output of cmd.
113        @raises CalledProcessError: If there was a non-0 return code.
114        """
115        out = subprocess.check_output(['googlesh', '-s', '-uchromeos-test',
116                                       '-m%s' % host, '%s' % cmd],
117                                      stderr=subprocess.STDOUT)
118        return out
119    return runner
120
121
122def execute_command(host, cmd, **kwargs):
123    """
124    Executes a command on the host `host`.  This an optimization that if
125    we're already chromeos-test, we can just ssh to the machine in question.
126    Or if we're local, we don't have to ssh at all.
127
128    @param host: The hostname to execute the command on.
129    @param cmd: The command to run.  Special shell syntax (such as pipes)
130                is allowed.
131    @param kwargs: Key word arguments for the runner functions.
132    @returns: The output of the command.
133    """
134    if utils.is_localhost(host):
135        runner = local_runner
136    elif getpass.getuser() == 'chromeos-test':
137        runner = host_object_runner(host)
138    else:
139        runner = googlesh_runner(host)
140
141    return runner(cmd, **kwargs)
142