1#!/usr/bin/python
2
3# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Module storing common utilities used by the deploy_autotest package.
8
9This python module is shared between the two programs in deploy_autotest.
10Specifically it contains shared constants and their shared arg parser.
11"""
12
13import argparse
14import logging
15
16def setup_logging():
17    """Setup basic logging with all logging info stripped."""
18    screen_handler = logging.StreamHandler()
19    screen_handler.setFormatter(logging.Formatter('%(message)s'))
20    logging.getLogger().addHandler(screen_handler)
21    logging.getLogger().setLevel(logging.INFO)
22
23
24SYNC = 'sync'
25RESTART = 'restart'
26PRINT = 'print'
27VALID_COMMANDS = [SYNC, RESTART, PRINT]
28
29DEVS = 'devservers'
30DRONES = 'drones'
31SCHEDULER = 'scheduler'
32VALID_TARGETS = [DEVS, DRONES, SCHEDULER]
33
34
35def parse_args(argv):
36    parser = argparse.ArgumentParser()
37    parser.add_argument('operation',
38                        help='Operation to perform. Must be one of: %s' %
39                        ' '.join(VALID_COMMANDS))
40    parser.add_argument('servers', nargs='+',
41                        help='Any set of items from the list: %s' %
42                        ' '.join(VALID_TARGETS))
43    parsed_args = parser.parse_args(argv)
44
45    # Some sanity checks.
46    if not parsed_args.operation in VALID_COMMANDS:
47        parser.error('Invalid operation specified. Must be one of: %s' %
48                     ' '.join(VALID_COMMANDS))
49
50    if not set(parsed_args.servers).issubset(set(VALID_TARGETS)):
51        parser.error('All servers must be one of %s' % ' '.join(VALID_TARGETS))
52
53    return parsed_args
54