1#  Copyright (C) 2023 The Android Open Source Project
2#
3#  Licensed under the Apache License, Version 2.0 (the "License");
4#  you may not use this file except in compliance with the License.
5#  You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#  Unless required by applicable law or agreed to in writing, software
10#  distributed under the License is distributed on an "AS IS" BASIS,
11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#  See the License for the specific language governing permissions and
13#  limitations under the License.
14
15import argparse
16import sys
17# TODO(b/311467339): add asbl support for test runs in the lab
18# from absl import flags
19from mobly import test_runner
20
21
22"""
23Get Test Args
24
25Extracts the test arguments from the System Args
26
27usage:
28  python3 <test> -- -c /tmp/config.yaml --test_args=k1=v1 --test_args=k2=v2
29
30CATBox usage:
31  catbox-tradefed run commandAndExit <test-plan> --mobly-options --test_args=k1=v1 --mobly-options --test_args=k2=v2
32
33Returns: Dictionary with key-value pair
34  e.g.
35  {
36    k1: v1,
37    k2: v2
38  }
39"""
40def get_test_args():
41    parser = argparse.ArgumentParser(description='Parse Test Args.')
42    group = parser.add_mutually_exclusive_group(required=False)
43    group.add_argument('--test_args',
44                          action='append',
45                          nargs='+',
46                          type=str,
47                          help='A list of test args for the test.')
48    parsed_test_args = parser.parse_known_args(sys.argv)[0]
49
50    if not parsed_test_args.test_args:
51        return {}
52
53    test_args_list = [
54        test_arg
55        for test_args in parsed_test_args.test_args
56        for test_arg in test_args
57    ]
58    test_args = {test_arg.split("=")[0]: test_arg.split("=")[1] for test_arg in test_args_list}
59    return test_args
60
61"""Pass test arguments after '--' to the test runner. Needed for Mobly Test Runner.
62
63Splits the arguments vector by '--'. Anything before separtor is treated as absl flags.
64Everything after is a Mobly Test Runner arguments. Example:
65
66    python3 <test> --itreations=2 -- -c /tmp/config.yaml
67
68Example usage:
69
70    if __name__ == '__main__':
71        common_main()
72"""
73def common_main():
74    absl_argv = []
75    if '--' in sys.argv:
76        index = sys.argv.index('--')
77        absl_argv = sys.argv[:index]
78        sys.argv = sys.argv[:1] + sys.argv[index + 1:]
79    # TODO(b/311467339): add asbl support for test runs in the lab
80    # flags.FLAGS(absl_argv)
81    test_runner.main()
82