Support selecting test cases in a class with regex (#938)

Users can now pass in regex test selector with prefix `re:` in the test case selection arg, for the test class level.

mobly/base_suite.py

@@ -47,8 +47,10 @@ def add_test_class(self, clazz, config=None, tests=None, name_suffix=None):
       config: config_parser.TestRunConfig, the config to run the class with. If
         not specified, the default config passed from google3 infra is used.
       tests: list of strings, names of the tests to run in this test class, in
-        the execution order. If not specified, all tests in the class are
-        executed.
+        the execution order. Or a string with prefix `re:` for full regex match
+        of test cases; all matched test cases will be executed; an error is
+        raised if no match is found.
+        If not specified, all tests in the class are executed.
       name_suffix: string, suffix to append to the class name for reporting.
         This is used for differentiating the same class executed with different
         parameters in a suite.

mobly/base_test.py

@@ -19,6 +19,7 @@
 import inspect
 import logging
 import os
+import re
 import sys
 
 from mobly import controller_manager
@@ -31,6 +32,7 @@
 # Macro strings for test result reporting.
 TEST_CASE_TOKEN = '[Test]'
 RESULT_LINE_TEMPLATE = TEST_CASE_TOKEN + ' %s %s'
+TEST_SELECTOR_REGEX_PREFIX = 're:'
 
 TEST_STAGE_BEGIN_LOG_TEMPLATE = '[{parent_token}]#{child_token} >>> BEGIN >>>'
 TEST_STAGE_END_LOG_TEMPLATE = '[{parent_token}]#{child_token} <<< END <<<'
@@ -1003,7 +1005,8 @@ def _get_test_methods(self, test_names):
     """Resolves test method names to bound test methods.
 
     Args:
-      test_names: A list of strings, each string is a test method name.
+      test_names: A list of strings, each string is a test method name or a
+        regex for matching test names.
 
     Returns:
       A list of tuples of (string, function). String is the test method
@@ -1014,21 +1017,52 @@ def _get_test_methods(self, test_names):
         This can only be caused by user input.
     """
     test_methods = []
+    # Process the test name selector one by one.
     for test_name in test_names:
-      if not test_name.startswith('test_'):
-        raise Error(
-            'Test method name %s does not follow naming '
-            'convention test_*, abort.' % test_name
+      if test_name.startswith(TEST_SELECTOR_REGEX_PREFIX):
+        # process the selector as a regex.
+        regex_matching_methods = self._get_regex_matching_test_methods(
+            test_name.removeprefix(TEST_SELECTOR_REGEX_PREFIX)
         )
+        test_methods += regex_matching_methods
+        continue
+      # process the selector as a regular test name string.
+      self._assert_valid_test_name(test_name)
+      if test_name not in self.get_existing_test_names():
+        raise Error(f'{self.TAG} does not have test method {test_name}.')
       if hasattr(self, test_name):
         test_method = getattr(self, test_name)
       elif test_name in self._generated_test_table:
         test_method = self._generated_test_table[test_name]
-      else:
-        raise Error('%s does not have test method %s.' % (self.TAG, test_name))
       test_methods.append((test_name, test_method))
     return test_methods
 
+  def _get_regex_matching_test_methods(self, test_name_regex):
+    matching_name_tuples = []
+    for name, method in inspect.getmembers(self, callable):
+      if (
+          name.startswith('test_')
+          and re.fullmatch(test_name_regex, name) is not None
+      ):
+        matching_name_tuples.append((name, method))
+    for name, method in self._generated_test_table.items():
+      if re.fullmatch(test_name_regex, name) is not None:
+        self._assert_valid_test_name(name)
+        matching_name_tuples.append((name, method))
+    if not matching_name_tuples:
+      raise Error(
+          f'{test_name_regex} does not match with any valid test case '
+          f'in {self.TAG}, abort!'
+      )
+    return matching_name_tuples
+
+  def _assert_valid_test_name(self, test_name):
+    if not test_name.startswith('test_'):
+      raise Error(
+          'Test method name %s does not follow naming '
+          'convention test_*, abort.' % test_name
+      )
+
   def _skip_remaining_tests(self, exception):
     """Marks any requested test that has not been executed in a class as
     skipped.

mobly/test_runner.py

@@ -126,8 +126,12 @@ def parse_mobly_cli_args(argv):
       '--test_case',
       nargs='+',
       type=str,
-      metavar='[test_a test_b...]',
-      help='A list of tests in the test class to execute.',
+      metavar='[test_a test_b re:test_(c|d)...]',
+      help=(
+          'A list of tests in the test class to execute. Each value can be a '
+          'test name string or a `re:` prefixed string for full regex match of'
+          ' test names.'
+      ),
   )
   parser.add_argument(
       '-tb',

tests/mobly/base_test_test.py

@@ -201,6 +201,96 @@ def test_never(self):
     actual_record = bt_cls.results.passed[0]
     self.assertEqual(actual_record.test_name, 'test_something')
 
+  def test_cli_test_selection_with_regex(self):
+    class MockBaseTest(base_test.BaseTestClass):
+
+      def __init__(self, controllers):
+        super().__init__(controllers)
+        self.tests = ('test_never',)
+
+      def test_foo(self):
+        pass
+
+      def test_a(self):
+        pass
+
+      def test_b(self):
+        pass
+
+      def test_something_1(self):
+        pass
+
+      def test_something_2(self):
+        pass
+
+      def test_something_3(self):
+        pass
+
+      def test_never(self):
+        # This should not execute since it's not selected by cmd line input.
+        never_call()
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    bt_cls.run(test_names=['re:test_something_.*', 'test_foo', 're:test_(a|b)'])
+    self.assertEqual(len(bt_cls.results.passed), 6)
+    self.assertEqual(bt_cls.results.passed[0].test_name, 'test_something_1')
+    self.assertEqual(bt_cls.results.passed[1].test_name, 'test_something_2')
+    self.assertEqual(bt_cls.results.passed[2].test_name, 'test_something_3')
+    self.assertEqual(bt_cls.results.passed[3].test_name, 'test_foo')
+    self.assertEqual(bt_cls.results.passed[4].test_name, 'test_a')
+    self.assertEqual(bt_cls.results.passed[5].test_name, 'test_b')
+
+  def test_cli_test_selection_with_regex_generated_tests(self):
+    class MockBaseTest(base_test.BaseTestClass):
+
+      def __init__(self, controllers):
+        super().__init__(controllers)
+        self.tests = ('test_never',)
+
+      def pre_run(self):
+        self.generate_tests(
+            test_logic=self.logic,
+            name_func=lambda i: f'test_something_{i}',
+            arg_sets=[(i + 1,) for i in range(3)],
+        )
+
+      def test_foo(self):
+        pass
+
+      def logic(self, _):
+        pass
+
+      def test_never(self):
+        # This should not execute since it's not selected by cmd line input.
+        never_call()
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    bt_cls.run(test_names=['re:test_something_.*', 'test_foo'])
+    self.assertEqual(len(bt_cls.results.passed), 4)
+    self.assertEqual(bt_cls.results.passed[0].test_name, 'test_something_1')
+    self.assertEqual(bt_cls.results.passed[1].test_name, 'test_something_2')
+    self.assertEqual(bt_cls.results.passed[2].test_name, 'test_something_3')
+    self.assertEqual(bt_cls.results.passed[3].test_name, 'test_foo')
+
+  def test_cli_test_selection_with_regex_fail_by_convention(self):
+    class MockBaseTest(base_test.BaseTestClass):
+
+      def __init__(self, controllers):
+        super().__init__(controllers)
+        self.tests = ('test_never',)
+
+      def test_something(self):
+        pass
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    expected_msg = (
+        r'not_a_test_something does not match with any valid test case in '
+        r'MockBaseTest, abort!'
+    )
+    with self.assertRaisesRegex(base_test.Error, expected_msg):
+      bt_cls.run(test_names=['re:not_a_test_something'])
+    self.assertEqual(len(bt_cls.results.passed), 0)
+
   def test_cli_test_selection_fail_by_convention(self):
     class MockBaseTest(base_test.BaseTestClass):