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
5import argparse
6import logging
7import os
8import sys
9
10import common
11from autotest_lib.client.common_lib import error
12from autotest_lib.server import hosts
13from autotest_lib.server import utils
14from autotest_lib.server.hosts import moblab_host
15from autotest_lib.server.hosts import ssh_host
16
17
18_LOGGING_FORMAT = '%(asctime)s - %(levelname)s - %(message)s'
19_TEST_LAUNCH_SCRIPT = 'brillo_test_launcher.py'
20
21# Running against a virtual machine has several intricacies that we need to
22# adjust for. Namely SSH requires the use of 'localhost' while HTTP requires
23# the use of '127.0.0.1'. Also because we are forwarding the ports from the VM
24# to the host system, the ports to use for these services are also different
25# from running on a physical machine.
26_VIRT_MACHINE_SSH_ADDR = 'localhost:9222'
27_VIRT_MACHINE_AFE_ADDR = '127.0.0.1:8888'
28_VIRT_MACHINE_DEVSERVER_PORT = '7777'
29_PHYS_MACHINE_DEVSERVER_PORT = '8080'
30_MOBLAB_MIN_VERSION = 7569
31_MOBLAB_IMAGE_DOWNLOAD_URL = ('https://storage.googleapis.com/chromeos-image-'
32                              'archive/moblab_brillo_images/'
33                              'moblab_brillo_%s.bin' % _MOBLAB_MIN_VERSION)
34
35
36class BrilloTestError(Exception):
37    """A general error while testing Brillo."""
38
39
40class BrilloMoblabInitializationError(BrilloTestError):
41    """An error during Moblab initialization or handling."""
42
43
44def get_moblab_and_devserver_port(moblab_hostname):
45    """Initializes and returns a MobLab Host Object.
46
47    @params moblab_hostname: The Moblab hostname, None if using a local virtual
48                             machine.
49
50    @returns A pair consisting of a MoblabHost and a devserver port.
51
52    @raise BrilloMoblabInitializationError: Failed to set up the Moblab.
53    """
54    if moblab_hostname:
55        web_address = moblab_hostname
56        devserver_port = _PHYS_MACHINE_DEVSERVER_PORT
57        rpc_timeout_min = 2
58    else:
59        moblab_hostname = _VIRT_MACHINE_SSH_ADDR
60        web_address = _VIRT_MACHINE_AFE_ADDR
61        devserver_port = _VIRT_MACHINE_DEVSERVER_PORT
62        rpc_timeout_min = 5
63
64    try:
65        host = hosts.create_host(moblab_hostname,
66                                 host_class=moblab_host.MoblabHost,
67                                 connectivity_class=ssh_host.SSHHost,
68                                 web_address=web_address,
69                                 retain_image_storage=True,
70                                 rpc_timeout_min=rpc_timeout_min)
71    except error.AutoservRunError as e:
72        raise BrilloMoblabInitializationError(
73                'Unable to connect to the MobLab: %s' % e)
74
75    moblab_version = int(host.get_release_version().split('.')[0])
76    if moblab_version < _MOBLAB_MIN_VERSION:
77        raise BrilloMoblabInitializationError(
78                'The Moblab version (%s) is older than the minimum required '
79                '(%s). Download a current version from URL: %s' %
80                (moblab_version, _MOBLAB_MIN_VERSION,
81                _MOBLAB_IMAGE_DOWNLOAD_URL))
82
83    try:
84        host.afe.get_hosts()
85    except Exception as e:
86        raise BrilloMoblabInitializationError(
87                "Unable to communicate with the MobLab's web frontend, "
88                "please verify that it is up and running at http://%s/\n"
89                "Error: %s" % (host.web_address, e))
90
91    return host, devserver_port
92
93
94def parse_args(description, setup_parser=None, validate_args=None):
95    """Parse command-line arguments.
96
97    @param description: The script description in the help message.
98    @param setup_parser: Function that takes a parser object and adds
99                         script-specific options to it.
100    @param validate_args: Function that takes a parser object and the parsed
101                          arguments and validates the arguments. It should use
102                          parser.error() to report errors.
103
104    @return Parsed and validated arguments.
105    """
106    parser = argparse.ArgumentParser(description=description)
107    if setup_parser:
108        setup_parser(parser)
109
110    # Add common options.
111    parser.add_argument('-m', '--moblab_host',
112                        help='MobLab hostname or IP to launch tests. If this '
113                             'argument is not provided, the test launcher '
114                             'will attempt to test a local virtual machine '
115                             'instance of MobLab.')
116    parser.add_argument('-a', '--adb_host',
117                        help='Hostname or IP of the adb_host connected to the '
118                             'Brillo DUT. Default is to assume it is connected '
119                             'directly to the MobLab.')
120    parser.add_argument('-n', '--no_quickmerge', dest='quickmerge',
121                        action='store_false',
122                        help='Do not update the Autotest code on the Moblab')
123    parser.add_argument('-d', '--debug', action='store_true',
124                        help='Print log statements.')
125
126    args = parser.parse_args()
127
128    # Configure the root logger.
129    logging.getLogger().setLevel(logging.DEBUG if args.debug else logging.INFO)
130    for log_handler in logging.getLogger().handlers:
131        log_handler.setFormatter(logging.Formatter(fmt=_LOGGING_FORMAT))
132
133    if validate_args:
134        validate_args(parser, args)
135
136    return args
137
138
139def setup_test_action_parser(parser):
140    """Add parser options related to test action.
141
142    @param parser: argparse.ArgumentParser of the script.
143    """
144    launch_opts = parser.add_mutually_exclusive_group()
145    launch_opts.add_argument('-A', '--print_args', action='store_true',
146                             help='Print test arguments to stdout instead of '
147                                  'launching the test.')
148    launch_opts.add_argument('-C', '--print_command', action='store_true',
149                             help='Print complete test launch command instead '
150                                  'of launching the test.')
151
152
153def _get_arg_strs(test_args):
154    """Converts an argument dictionary into a list of 'arg=val' strings."""
155    return ['%s=%s' % kv for kv in test_args.iteritems()]
156
157
158def _get_command(moblab, test_name, test_args, do_quickmerge, do_quote):
159    """Returns the test launch command.
160
161    @param moblab: MoblabHost representing the MobLab being used for testing.
162    @param test_name: The name of the test to run.
163    @param test_args: Dictionary of test arguments.
164    @param do_quickmerge: If False, pass the --no-quickmerge flag.
165    @param do_quote: If True, add single-quotes around test arguments.
166
167    @return Test launch command as a list of strings.
168    """
169    # pylint: disable=missing-docstring
170    def quote(val):
171        return "'%s'" % val if do_quote else val
172
173    cmd = [os.path.join(os.path.dirname(__file__), _TEST_LAUNCH_SCRIPT),
174           '-t', quote(test_name)]
175    if not do_quickmerge:
176        cmd.append('-n')
177    if not moblab.hostname.startswith('localhost'):
178           cmd += ['-m', quote(moblab.hostname)]
179    for arg_str in _get_arg_strs(test_args):
180        cmd += ['-A', quote(arg_str)]
181    return cmd
182
183
184def _print_args(test_args):
185    """Prints the test arguments to stdout, one per line.
186
187    @param test_args: Dictionary of test arguments.
188    """
189    print('\n'.join(_get_arg_strs(test_args)))
190
191
192def _print_command(moblab, test_name, test_args, do_quickmerge):
193    """Prints the test launch command to stdout with quoting.
194
195    @param moblab: MoblabHost representing the MobLab being used for testing.
196    @param test_name: The name of the test to run.
197    @param test_args: Dictionary of test arguments.
198    @param do_quickmerge: If False, pass the --no-quickmerge flag.
199    """
200    print(' '.join(
201            _get_command(moblab, test_name, test_args, do_quickmerge, True)))
202
203
204def _run_command(moblab, test_name, test_args, do_quickmerge):
205    """Runs the test launch script.
206
207    @param moblab: MoblabHost representing the MobLab being used for testing.
208    @param test_name: The name of the test to run.
209    @param test_args: Dictionary of test arguments.
210    @param do_quickmerge: If False, pass the --no_quickmerge flag.
211    """
212    utils.run(_get_command(moblab, test_name, test_args, do_quickmerge, False),
213              stdout_tee=sys.stdout, stderr_tee=sys.stderr)
214
215
216def do_test_action(args, moblab, test_name, test_args):
217    """Performs the desired action related to the test.
218
219    @param args: Parsed arguments.
220    @param moblab: MoblabHost representing the MobLab being used for testing.
221    @param test_name: The name of the test to run.
222    @param test_args: Dictionary of test arguments.
223    """
224    if args.print_args:
225        logging.info('Printing test arguments')
226        _print_args(test_args)
227    elif args.print_command:
228        logging.info('Printing test launch command')
229        _print_command(moblab, test_name, test_args, args.quickmerge)
230    else:
231        logging.info('Launching test')
232        _run_command(moblab, test_name, test_args, args.quickmerge)
233