Home | History | Annotate | Download | only in site_utils
      1 #!/usr/bin/python
      2 # Copyright 2015 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 import logging
      7 import os
      8 import sys
      9 
     10 import common
     11 try:
     12     # Ensure the chromite site-package is installed.
     13     from chromite.lib import *
     14 except ImportError:
     15     import subprocess
     16     build_externals_path = os.path.join(
     17             os.path.dirname(os.path.dirname(os.path.realpath(__file__))),
     18             'utils', 'build_externals.py')
     19     subprocess.check_call([build_externals_path, 'chromiterepo'])
     20     # Restart the script so python now finds the autotest site-packages.
     21     sys.exit(os.execv(__file__, sys.argv))
     22 from autotest_lib.site_utils import brillo_common
     23 
     24 
     25 _TEST_NAME = 'brillo_RecoverFromBadImage'
     26 
     27 
     28 def setup_parser(parser):
     29     """Add parser options.
     30 
     31     @param parser: argparse.ArgumentParser of the script.
     32     """
     33     parser.add_argument('-i', '--recovery_image', metavar='FILE', required=True,
     34                         help='Image file to use for recovery. This is a '
     35                              'mandatory input.')
     36     parser.add_argument('-p', '--partition', metavar='NAME',
     37                         help='Name of partition to recover. If the name ends '
     38                              'with "_X" then it will be substitued with the '
     39                              'currently active slot (e.g. "_a"). (default: '
     40                              'system_X)')
     41     parser.add_argument('-D', '--device', metavar='PATH',
     42                         help='Path of partition device. (default: infer from '
     43                              'name)')
     44 
     45     brillo_common.setup_test_action_parser(parser)
     46 
     47 
     48 def main(args):
     49     """The main function."""
     50     args = brillo_common.parse_args(
     51             'Set up Moblab for running Brillo image recovery test, then launch '
     52             'the test (unless otherwise requested).',
     53             setup_parser=setup_parser)
     54 
     55     test_args = {}
     56     if args.partition:
     57         test_args['partition'] = args.partition
     58     if args.device:
     59         test_args['device'] = args.device
     60 
     61     moblab, _ = brillo_common.get_moblab_and_devserver_port(args.moblab_host)
     62     tmp_dir = moblab.make_tmp_dir()
     63     try:
     64         remote_recovery_image = os.path.join(
     65                 tmp_dir, os.path.basename(args.recovery_image))
     66         moblab.send_file(args.recovery_image, remote_recovery_image)
     67         moblab.run('chown -R moblab:moblab %s' % tmp_dir)
     68         test_args['image_file'] = remote_recovery_image
     69         logging.info('Recovery image was staged')
     70         brillo_common.do_test_action(args, moblab, _TEST_NAME, test_args)
     71     finally:
     72         moblab.run('rm -rf %s' % tmp_dir)
     73 
     74 
     75 if __name__ == '__main__':
     76     try:
     77         main(sys.argv)
     78         sys.exit(0)
     79     except brillo_common.BrilloTestError as e:
     80         logging.error('Error: %s', e)
     81 
     82     sys.exit(1)
     83