1# Copyright 2018 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"""Wrapper for running suites of tests and waiting for completion."""
6
7from __future__ import absolute_import
8from __future__ import division
9from __future__ import print_function
10
11import os
12import sys
13
14import logging
15
16from lucifer import autotest
17from skylab_suite import cros_suite
18from skylab_suite import suite_parser
19from skylab_suite import suite_runner
20from skylab_suite import suite_tracking
21from skylab_suite import swarming_lib
22
23
24PROVISION_SUITE_NAME = 'provision'
25
26
27def _parse_suite_handler_spec(options):
28    provision_num_required = 0
29    if 'num_required' in options.suite_args:
30        provision_num_required = options.suite_args['num_required']
31
32    return cros_suite.SuiteHandlerSpec(
33            suite_name=options.suite_name,
34            wait=not options.create_and_return,
35            suite_id=options.suite_id,
36            timeout_mins=options.timeout_mins,
37            passed_mins=options.passed_mins,
38            test_retry=options.test_retry,
39            max_retries=options.max_retries,
40            provision_num_required=provision_num_required)
41
42
43def _should_run(suite_spec):
44    tags = {'build': suite_spec.test_source_build,
45            'suite': suite_spec.suite_name}
46    tasks = swarming_lib.query_task_by_tags(tags)
47    current_task_id = suite_tracking.get_task_id_for_task_summaries(
48            os.environ.get('SWARMING_TASK_ID'))
49    logging.info('The current task id is: %s', current_task_id)
50    extra_task_ids = set([])
51    for t in tasks:
52        if t['task_id'] != current_task_id:
53            extra_task_ids.add(t['task_id'])
54
55    return extra_task_ids
56
57
58def _run_suite(options):
59    swarming_client = swarming_lib.Client(options.swarming_auth_json)
60    run_suite_common = autotest.load('site_utils.run_suite_common')
61    logging.info('Kicked off suite %s', options.suite_name)
62    suite_spec = suite_parser.parse_suite_spec(options)
63    if options.pre_check:
64        extra_task_ids = _should_run(suite_spec)
65        if extra_task_ids:
66            logging.info(
67                    'The same suites are already run in the past: \n%s',
68                    '\n'.join([swarming_lib.get_task_link(tid)
69                               for tid in extra_task_ids]))
70            return run_suite_common.SuiteResult(
71                    run_suite_common.RETURN_CODES.OK)
72
73    if options.suite_name == PROVISION_SUITE_NAME:
74        suite_job = cros_suite.ProvisionSuite(suite_spec, swarming_client)
75    else:
76        suite_job = cros_suite.Suite(suite_spec, swarming_client)
77
78    try:
79        suite_job.prepare()
80    except Exception as e:
81        logging.exception('Infra failure in setting up suite job')
82        return run_suite_common.SuiteResult(
83                run_suite_common.RETURN_CODES.INFRA_FAILURE)
84
85    suite_handler_spec = _parse_suite_handler_spec(options)
86    suite_handler = cros_suite.SuiteHandler(suite_handler_spec, swarming_client)
87    suite_runner.run(swarming_client,
88                     suite_job.test_specs,
89                     suite_handler,
90                     options.dry_run)
91
92    if options.create_and_return:
93        suite_tracking.log_create_task(suite_job.suite_name,
94                                       suite_handler.suite_id)
95        suite_tracking.print_child_test_annotations(suite_handler)
96        return run_suite_common.SuiteResult(run_suite_common.RETURN_CODES.OK)
97
98    return_code = suite_tracking.log_suite_results(
99                suite_job.suite_name, suite_handler)
100    return run_suite_common.SuiteResult(return_code)
101
102
103def parse_args():
104    """Parse & validate skylab suite args."""
105    parser = suite_parser.make_parser()
106    options = parser.parse_args()
107    if options.do_nothing:
108        logging.info('Exit early because --do_nothing requested.')
109        sys.exit(0)
110
111    if not suite_parser.verify_and_clean_options(options):
112        parser.print_help()
113        sys.exit(1)
114
115    return options
116
117
118def _setup_env(options):
119    """Set environment variables based on commandline options."""
120    os.environ['SWARMING_CREDS'] = options.swarming_auth_json
121
122
123def main():
124    """Entry point."""
125    autotest.monkeypatch()
126
127    options = parse_args()
128    _setup_env(options)
129    suite_tracking.setup_logging()
130    result = _run_suite(options)
131    logging.info('Will return from %s with status: %s',
132                 os.path.basename(__file__), result.string_code)
133    return result.return_code
134
135
136if __name__ == "__main__":
137    sys.exit(main())
138