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
5import argparse
6import getpass
7import logging
8import sys
9
10import common
11from autotest_lib.client.common_lib import utils
12
13
14def 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
33def 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
52class 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.
83config = Config()
84
85
86def 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