Home | History | Annotate | Download | only in lxc 
      1 # Copyright 2017 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 
      5 import argparse
      6 import getpass
      7 import logging
      8 import sys
      9 
     10 import common
     11 from autotest_lib.client.common_lib import utils
     12 
     13 
     14 def setup_logging(log_level):
     15     """Sets up direct logging to stdout for unittests.
     16 
     17     @param log_level: Level of logging to redirect to stdout, default to INFO.
     18     """
     19     # Lifted from client.common_lib.logging_config.
     20     FORMAT = ('%(asctime)s.%(msecs)03d %(levelname)-5.5s|%(module)18.18s:'
     21               '%(lineno)4.4d| %(threadName)16.16s(%(thread)d)| %(message)s')
     22 
     23     logger = logging.getLogger()
     24     logger.setLevel(log_level)
     25     handler = logging.StreamHandler(sys.stdout)
     26     handler.setLevel(log_level)
     27     formatter = logging.Formatter(FORMAT)
     28     handler.setFormatter(formatter)
     29     logger.handlers = []
     30     logger.addHandler(handler)
     31 
     32 
     33 def verify_user(require_sudo=True):
     34     """Checks that the current user is not root, but has sudo.
     35 
     36     Running unit tests as root can mask permissions problems, as not all the
     37     code runs as root in production.
     38     """
     39     # Ensure this process is not running as root.
     40     if getpass.getuser() == 'root':
     41         raise EnvironmentError('Unittests should not be run as root.')
     42 
     43     # However, most of the unit tests do require sudo.
     44     # TODO(dshi): crbug.com/459344 Set remove this enforcement when test
     45     # container can be unprivileged container.
     46     if require_sudo and utils.sudo_require_password():
     47         logging.warn('SSP requires root privilege to run commands, please '
     48                      'grant root access to this process.')
     49         utils.run('sudo true')
     50 
     51 
     52 class Config(object):
     53     """A class for parsing and storing command line options.
     54 
     55     A convenience class for helping with unit test setup.  A global instance of
     56     this class is set up by the setup function.  Clients can then check this
     57     object for flags set on the command line.
     58     """
     59     def parse_options(self):
     60         """Parses command line flags for unittests."""
     61         parser = argparse.ArgumentParser()
     62         parser.add_argument('-v', '--verbose', action='store_true',
     63                             help='Print out ALL entries.')
     64         parser.add_argument('-s', '--skip_cleanup', action='store_true',
     65                             help='Skip deleting test containers.')
     66         args, argv = parser.parse_known_args()
     67 
     68         for attr, value in vars(args).items():
     69             setattr(self, attr, value)
     70 
     71         # Hack: python unittest also processes args.  Construct an argv to pass
     72         # to it, that filters out the options it won't recognize.  Then replace
     73         # sys.argv with the constructed argv so that calling unittest.main "just
     74         # works".
     75         if args.verbose:
     76             argv.insert(0, '-v')
     77         argv.insert(0, sys.argv[0])
     78         sys.argv = argv
     79 
     80 
     81 # Global namespace object for storing unittest options specified on the command
     82 # line.
     83 config = Config()
     84 
     85 
     86 def setup(require_sudo=True):
     87     """Performs global setup for unit-tests."""
     88     config.parse_options()
     89 
     90     verify_user(require_sudo)
     91 
     92     log_level = logging.DEBUG if config.verbose else logging.INFO
     93     setup_logging(log_level)
     94