1# Copyright 2017 Google Inc.
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"""Runner for Mobly test suites.
15
16These is just example code to help users run a collection of Mobly test
17classes. Users can use it as is or customize it based on their requirements.
18
19There are two ways to use this runner.
20
211. Call suite_runner.run_suite() with one or more individual test classes. This
22is for users who just need to execute a collection of test classes without any
23additional steps.
24
25.. code-block:: python
26
27  from mobly import suite_runner
28
29  from my.test.lib import foo_test
30  from my.test.lib import bar_test
31  ...
32  if __name__ == '__main__':
33    suite_runner.run_suite(foo_test.FooTest, bar_test.BarTest)
34
352. Create a subclass of base_suite.BaseSuite and add the individual test
36classes. Using the BaseSuite class allows users to define their own setup
37and teardown steps on the suite level as well as custom config for each test
38class.
39
40.. code-block:: python
41
42  from mobly import base_suite
43  from mobly import suite_runner
44
45  from my.path import MyFooTest
46  from my.path import MyBarTest
47
48
49  class MySuite(base_suite.BaseSuite):
50
51    def setup_suite(self, config):
52      # Add a class with default config.
53      self.add_test_class(MyFooTest)
54      # Add a class with test selection.
55      self.add_test_class(MyBarTest,
56                          tests=['test_a', 'test_b'])
57      # Add the same class again with a custom config and suffix.
58      my_config = some_config_logic(config)
59      self.add_test_class(MyBarTest,
60                          config=my_config,
61                          name_suffix='WithCustomConfig')
62
63
64  if __name__ == '__main__':
65    suite_runner.run_suite_class()
66"""
67import argparse
68import collections
69import inspect
70import logging
71import sys
72
73from mobly import base_test
74from mobly import base_suite
75from mobly import config_parser
76from mobly import signals
77from mobly import test_runner
78
79
80class Error(Exception):
81  pass
82
83
84def _parse_cli_args(argv):
85  """Parses cli args that are consumed by Mobly.
86
87  Args:
88    argv: A list that is then parsed as cli args. If None, defaults to cli
89      input.
90
91  Returns:
92    Namespace containing the parsed args.
93  """
94  parser = argparse.ArgumentParser(description='Mobly Suite Executable.')
95  group = parser.add_mutually_exclusive_group(required=True)
96  group.add_argument(
97      '-c',
98      '--config',
99      type=str,
100      metavar='<PATH>',
101      help='Path to the test configuration file.',
102  )
103  group.add_argument(
104      '-l',
105      '--list_tests',
106      action='store_true',
107      help=(
108          'Print the names of the tests defined in a script without '
109          'executing them.'
110      ),
111  )
112  parser.add_argument(
113      '--tests',
114      '--test_case',
115      nargs='+',
116      type=str,
117      metavar='[ClassA[.test_a] ClassB[.test_b] ...]',
118      help='A list of test classes and optional tests to execute.',
119  )
120  parser.add_argument(
121      '-tb',
122      '--test_bed',
123      nargs='+',
124      type=str,
125      metavar='[<TEST BED NAME1> <TEST BED NAME2> ...]',
126      help='Specify which test beds to run tests on.',
127  )
128
129  parser.add_argument(
130      '-v',
131      '--verbose',
132      action='store_true',
133      help='Set console logger level to DEBUG',
134  )
135  if not argv:
136    argv = sys.argv[1:]
137  return parser.parse_known_args(argv)[0]
138
139
140def _find_suite_class():
141  """Finds the test suite class in the current module.
142
143  Walk through module members and find the subclass of BaseSuite. Only
144  one subclass is allowed in a module.
145
146  Returns:
147      The test suite class in the test module.
148  """
149  test_suites = []
150  main_module_members = sys.modules['__main__']
151  for _, module_member in main_module_members.__dict__.items():
152    if inspect.isclass(module_member):
153      if issubclass(module_member, base_suite.BaseSuite):
154        test_suites.append(module_member)
155  if len(test_suites) != 1:
156    logging.error(
157        'Expected 1 test class per file, found %s.',
158        [t.__name__ for t in test_suites],
159    )
160    sys.exit(1)
161  return test_suites[0]
162
163
164def _print_test_names(test_classes):
165  """Prints the names of all the tests in all test classes.
166  Args:
167    test_classes: classes, the test classes to print names from.
168  """
169  for test_class in test_classes:
170    cls = test_class(config_parser.TestRunConfig())
171    test_names = []
172    try:
173      # Executes pre-setup procedures, this is required since it might
174      # generate test methods that we want to return as well.
175      cls._pre_run()
176      if cls.tests:
177        # Specified by run list in class.
178        test_names = list(cls.tests)
179      else:
180        # No test method specified by user, list all in test class.
181        test_names = cls.get_existing_test_names()
182    except Exception:
183      logging.exception('Failed to retrieve generated tests.')
184    finally:
185      cls._clean_up()
186    print('==========> %s <==========' % cls.TAG)
187    for name in test_names:
188      print(f'{cls.TAG}.{name}')
189
190
191def run_suite_class(argv=None):
192  """Executes tests in the test suite.
193
194  Args:
195    argv: A list that is then parsed as CLI args. If None, defaults to sys.argv.
196  """
197  cli_args = _parse_cli_args(argv)
198  suite_class = _find_suite_class()
199  if cli_args.list_tests:
200    _print_test_names([suite_class])
201    sys.exit(0)
202  test_configs = config_parser.load_test_config_file(
203      cli_args.config, cli_args.test_bed
204  )
205  config_count = len(test_configs)
206  if config_count != 1:
207    logging.error('Expect exactly one test config, found %d', config_count)
208  config = test_configs[0]
209  runner = test_runner.TestRunner(
210      log_dir=config.log_path, testbed_name=config.testbed_name
211  )
212  suite = suite_class(runner, config)
213  console_level = logging.DEBUG if cli_args.verbose else logging.INFO
214  ok = False
215  with runner.mobly_logger(console_level=console_level):
216    try:
217      suite.setup_suite(config.copy())
218      try:
219        runner.run()
220        ok = runner.results.is_all_pass
221        print(ok)
222      except signals.TestAbortAll:
223        pass
224    finally:
225      suite.teardown_suite()
226  if not ok:
227    sys.exit(1)
228
229
230def run_suite(test_classes, argv=None):
231  """Executes multiple test classes as a suite.
232
233  This is the default entry point for running a test suite script file
234  directly.
235
236  Args:
237    test_classes: List of python classes containing Mobly tests.
238    argv: A list that is then parsed as cli args. If None, defaults to cli
239      input.
240  """
241  args = _parse_cli_args(argv)
242
243  # Check the classes that were passed in
244  for test_class in test_classes:
245    if not issubclass(test_class, base_test.BaseTestClass):
246      logging.error(
247          'Test class %s does not extend mobly.base_test.BaseTestClass',
248          test_class,
249      )
250      sys.exit(1)
251
252  if args.list_tests:
253    _print_test_names(test_classes)
254    sys.exit(0)
255
256  # Load test config file.
257  test_configs = config_parser.load_test_config_file(args.config, args.test_bed)
258  # Find the full list of tests to execute
259  selected_tests = compute_selected_tests(test_classes, args.tests)
260
261  console_level = logging.DEBUG if args.verbose else logging.INFO
262  # Execute the suite
263  ok = True
264  for config in test_configs:
265    runner = test_runner.TestRunner(config.log_path, config.testbed_name)
266    with runner.mobly_logger(console_level=console_level):
267      for test_class, tests in selected_tests.items():
268        runner.add_test_class(config, test_class, tests)
269      try:
270        runner.run()
271        ok = runner.results.is_all_pass and ok
272      except signals.TestAbortAll:
273        pass
274      except Exception:
275        logging.exception('Exception when executing %s.', config.testbed_name)
276        ok = False
277  if not ok:
278    sys.exit(1)
279
280
281def compute_selected_tests(test_classes, selected_tests):
282  """Computes tests to run for each class from selector strings.
283
284  This function transforms a list of selector strings (such as FooTest or
285  FooTest.test_method_a) to a dict where keys are test_name classes, and
286  values are lists of selected tests in those classes. None means all tests in
287  that class are selected.
288
289  Args:
290    test_classes: list of strings, names of all the classes that are part
291      of a suite.
292    selected_tests: list of strings, list of tests to execute. If empty,
293      all classes `test_classes` are selected. E.g.
294
295      .. code-block:: python
296
297        [
298          'FooTest',
299          'BarTest',
300          'BazTest.test_method_a',
301          'BazTest.test_method_b'
302        ]
303
304  Returns:
305    dict: Identifiers for TestRunner. Keys are test class names; valures
306      are lists of test names within class. E.g. the example in
307      `selected_tests` would translate to:
308
309      .. code-block:: python
310
311        {
312          FooTest: None,
313          BarTest: None,
314          BazTest: ['test_method_a', 'test_method_b']
315        }
316
317      This dict is easy to consume for `TestRunner`.
318  """
319  class_to_tests = collections.OrderedDict()
320  if not selected_tests:
321    # No selection is needed; simply run all tests in all classes.
322    for test_class in test_classes:
323      class_to_tests[test_class] = None
324    return class_to_tests
325
326  # The user is selecting some tests to run. Parse the selectors.
327  # Dict from test_name class name to list of tests to execute (or None for all
328  # tests).
329  test_class_name_to_tests = collections.OrderedDict()
330  for test_name in selected_tests:
331    if '.' in test_name:  # Has a test method
332      (test_class_name, test_name) = test_name.split('.', maxsplit=1)
333      if test_class_name not in test_class_name_to_tests:
334        # Never seen this class before
335        test_class_name_to_tests[test_class_name] = [test_name]
336      elif test_class_name_to_tests[test_class_name] is None:
337        # Already running all tests in this class, so ignore this extra
338        # test.
339        pass
340      else:
341        test_class_name_to_tests[test_class_name].append(test_name)
342    else:  # No test method; run all tests in this class.
343      test_class_name_to_tests[test_name] = None
344
345  # Now transform class names to class objects.
346  # Dict from test_name class name to instance.
347  class_name_to_class = {cls.__name__: cls for cls in test_classes}
348  for test_class_name, tests in test_class_name_to_tests.items():
349    test_class = class_name_to_class.get(test_class_name)
350    if not test_class:
351      raise Error('Unknown test_name class %s' % test_class_name)
352    class_to_tests[test_class] = tests
353
354  return class_to_tests
355