xref: /aosp_15_r20/external/autotest/site_utils/deployment/prepare/main.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1#!/usr/bin/python3 -u
2# Copyright 2019 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Tool to (re)prepare a DUT for lab deployment."""
7
8from __future__ import absolute_import
9from __future__ import division
10from __future__ import print_function
11
12import argparse
13import errno
14import logging
15import logging.config
16import os
17import sys
18
19import common
20from autotest_lib.client.common_lib import autotest_enum
21from autotest_lib.client.common_lib import logging_manager
22from autotest_lib.server import afe_utils
23from autotest_lib.server import server_logging_config
24from autotest_lib.server.hosts import file_store
25from autotest_lib.site_utils.deployment.prepare import dut as preparedut
26from autotest_lib.server.hosts import factory
27from autotest_lib.site_utils.admin_audit import rpm_validator
28
29
30RETURN_CODES = autotest_enum.AutotestEnum(
31        'OK',
32        'SERVO_VERIFICATION_FAILURE',
33        'STAGE_USB_FAILURE',
34        'INSTALL_FIRMWARE_FAILURE',
35        'INSTALL_TEST_IMAGE_FAILURE',
36        'PRE_DEPLOY_VERIFICATION_FAILURE',
37        'BOOT_FROM_RECOVERY_MODE_FAILURE',
38        'SETUP_LABSTATION_FAILURE',
39        'UPDATE_LABEL_FAILURE',
40        'OTHER_FAILURES',
41)
42
43_SERVO_UART_LOGS = 'servo_uart'
44
45
46class DutPreparationError(Exception):
47    """Generic error raised during DUT preparation."""
48
49
50def main():
51    """Tool to (re)prepare a DUT for lab deployment."""
52    opts = _parse_args()
53
54    # Create logging setting
55    logging_manager.configure_logging(
56            server_logging_config.ServerLoggingConfig(),
57            results_dir=opts.results_dir)
58
59    try:
60        host_info = _read_store(opts.host_info_file)
61    except Exception as err:
62        logging.error("fail to prepare: %s", err)
63        return RETURN_CODES.OTHER_FAILURES
64
65    with create_host(opts.hostname, host_info, opts.results_dir) as host:
66        if opts.dry_run:
67            logging.info('DRY RUN: Would have run actions %s', opts.actions)
68            return
69
70        is_labstation = (host_info.get().os == "labstation")
71
72        if 'servo-verification' in opts.actions:
73            try:
74                if not is_labstation:
75                    preparedut.verify_servo(host)
76            except Exception as err:
77                logging.error("fail to check servo: %s", err)
78                return RETURN_CODES.SERVO_VERIFICATION_FAILURE
79
80        if 'stage-usb' in opts.actions:
81            try:
82                repair_image = afe_utils.get_stable_cros_image_name_v2(
83                        host_info.get())
84                logging.info('Using repair image %s, obtained from AFE',
85                             repair_image)
86                preparedut.download_image_to_servo_usb(host, repair_image)
87            except Exception as err:
88                logging.error("fail to stage image to usb: %s", err)
89                return RETURN_CODES.STAGE_USB_FAILURE
90
91        if 'install-test-image' in opts.actions:
92            try:
93                preparedut.install_test_image(host)
94            except Exception as err:
95                logging.error("fail to install test image: %s", err)
96                return RETURN_CODES.INSTALL_TEST_IMAGE_FAILURE
97
98        if 'install-firmware' in opts.actions:
99            try:
100                preparedut.install_firmware(host)
101            except Exception as err:
102                logging.error("fail to install firmware: %s", err)
103                return RETURN_CODES.INSTALL_FIRMWARE_FAILURE
104
105        if 'verify-recovery-mode' in opts.actions:
106            try:
107                preparedut.verify_boot_into_rec_mode(host)
108            except Exception as err:
109                logging.error("fail to boot from recovery mode: %s", err)
110                return RETURN_CODES.BOOT_FROM_RECOVERY_MODE_FAILURE
111
112        # TODO (otabek): mix this step with update-label later.
113        if 'setup-labstation' in opts.actions:
114            try:
115                preparedut.setup_hwid_and_serialnumber(host)
116            except Exception as err:
117                logging.error("fail to setup labstation: %s", err)
118                return RETURN_CODES.SETUP_LABSTATION_FAILURE
119
120        if 'update-label' in opts.actions:
121            try:
122                preparedut.setup_hwid_and_serialnumber(host)
123                if not is_labstation:
124                    host.labels.update_labels(host, task_name='deploy')
125            except Exception as err:
126                logging.error("fail to update label: %s", err)
127                return RETURN_CODES.UPDATE_LABEL_FAILURE
128
129        if 'run-pre-deploy-verification' in opts.actions:
130            try:
131                if is_labstation:
132                    logging.info("testing RPM information on labstation")
133                    preparedut.verify_labstation_RPM_config_unsafe(host)
134                else:
135                    preparedut.verify_servo(host)
136                    preparedut.verify_battery_status(host)
137                    preparedut.verify_ccd_testlab_enable(host)
138                    rpm_validator.verify_unsafe(host)
139            except Exception as err:
140                logging.error("fail on pre-deploy verification: %s", err)
141                return RETURN_CODES.PRE_DEPLOY_VERIFICATION_FAILURE
142
143    return RETURN_CODES.OK
144
145
146def _parse_args():
147    parser = argparse.ArgumentParser(
148            description='Prepare / validate DUT for lab deployment.')
149
150    parser.add_argument(
151            'actions',
152            nargs='+',
153            choices=[
154                    'servo-verification', 'stage-usb', 'install-test-image',
155                    'install-firmware', 'verify-recovery-mode',
156                    'run-pre-deploy-verification', 'update-label',
157                    'setup-labstation'
158            ],
159            help='DUT preparation actions to execute.',
160    )
161    parser.add_argument(
162            '--dry-run',
163            action='store_true',
164            default=False,
165            help='Run in dry-run mode. No changes will be made to the DUT.',
166    )
167    parser.add_argument(
168            '--results-dir',
169            required=True,
170            help='Directory to drop logs and output artifacts in.',
171    )
172
173    parser.add_argument(
174            '--hostname',
175            required=True,
176            help='Hostname of the DUT to prepare.',
177    )
178    parser.add_argument(
179            '--host-info-file',
180            required=True,
181            help=('Full path to HostInfo file.'
182                  ' DUT inventory information is read from the HostInfo file.'
183                  ),
184    )
185
186    return parser.parse_args()
187
188
189def _read_store(path):
190    """Read a HostInfo from a file at path."""
191    store = file_store.FileStore(path)
192    return store
193
194
195def create_host(hostname, host_info, results_dir):
196    """Yield a hosts.CrosHost object with the given inventory information.
197
198    @param hostname: Hostname of the DUT.
199    @param info: A HostInfo with the inventory information to use.
200    @param results_dir: Path to directory for logs / output artifacts.
201
202    @yield server.hosts.CrosHost object.
203    """
204    info = host_info.get()
205    if not info.board:
206        raise DutPreparationError('No board in DUT labels')
207    if not info.model:
208        raise DutPreparationError('No model in DUT labels')
209
210    need_servo = info.os != 'labstation'
211    dut_logs_dir = None
212
213    if need_servo:
214        # We assume target host is a cros DUT by default
215        if 'servo_host' not in info.attributes:
216            raise DutPreparationError('No servo_host in DUT attributes')
217        if 'servo_port' not in info.attributes:
218            raise DutPreparationError('No servo_port in DUT attributes')
219
220        dut_logs_dir = os.path.join(results_dir, _SERVO_UART_LOGS)
221        try:
222            os.makedirs(dut_logs_dir)
223        except OSError as e:
224            if e.errno != errno.EEXIST:
225                raise
226
227    return factory.create_target_host(hostname,
228                                      host_info_store=host_info,
229                                      try_lab_servo=need_servo,
230                                      try_servo_repair=need_servo,
231                                      try_servo_recovery=need_servo,
232                                      servo_uart_logs_dir=dut_logs_dir)
233
234
235if __name__ == '__main__':
236    sys.exit(main())
237