Support test case selection when using suite class. (#947)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index f4399ac..97b0f8a 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -14,6 +14,8 @@
 
 import abc
 
+import logging
+
 
 class BaseSuite(abc.ABC):
   """Class used to define a Mobly suite.
@@ -34,11 +36,20 @@
   def __init__(self, runner, config):
     self._runner = runner
     self._config = config.copy()
+    self._test_selector = None
 
   @property
   def user_params(self):
     return self._config.user_params
 
+  def set_test_selector(self, test_selector):
+    """Sets test selector.
+
+    Don't override or call this method. This should only be used by the Mobly
+    framework.
+    """
+    self._test_selector = test_selector
+
   def add_test_class(self, clazz, config=None, tests=None, name_suffix=None):
     """Adds a test class to the suite.
 
@@ -51,10 +62,20 @@
         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.
+        CLI argument `tests` takes precedence over this argument.
       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.
     """
+    if self._test_selector:
+      cls_name = clazz.__name__
+      if cls_name not in self._test_selector:
+        logging.info(
+            'Skipping test class %s due to CLI argument `tests`.', cls_name
+        )
+        return
+      tests = self._test_selector[cls_name]
+
     if not config:
       config = self._config
     self._runner.add_test_class(config, clazz, tests, name_suffix)
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index b0b064c..3dd6c29 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -210,6 +210,8 @@
       log_dir=config.log_path, testbed_name=config.testbed_name
   )
   suite = suite_class(runner, config)
+  test_selector = _parse_raw_test_selector(cli_args.tests)
+  suite.set_test_selector(test_selector)
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
   with runner.mobly_logger(console_level=console_level):
@@ -287,8 +289,8 @@
   that class are selected.
 
   Args:
-    test_classes: list of strings, names of all the classes that are part
-      of a suite.
+    test_classes: list of `type[base_test.BaseTestClass]`, all the test classes
+      that are part of a suite.
     selected_tests: list of strings, list of tests to execute. If empty,
       all classes `test_classes` are selected. E.g.
 
@@ -326,6 +328,50 @@
   # The user is selecting some tests to run. Parse the selectors.
   # Dict from test_name class name to list of tests to execute (or None for all
   # tests).
+  test_class_name_to_tests = _parse_raw_test_selector(selected_tests)
+
+  # Now compute the tests to run for each test class.
+  # Dict from test class name to class instance.
+  class_name_to_class = {cls.__name__: cls for cls in test_classes}
+  for test_class_name, tests in test_class_name_to_tests.items():
+    test_class = class_name_to_class.get(test_class_name)
+    if not test_class:
+      raise Error('Unknown test_class name %s' % test_class_name)
+    class_to_tests[test_class] = tests
+
+  return class_to_tests
+
+
+def _parse_raw_test_selector(selected_tests):
+  """Parses test selector from CLI arguments.
+
+  This function transforms a list of selector strings (such as FooTest or
+  FooTest.test_method_a) to a dict where keys are test_name classes, and
+  values are lists of selected tests in those classes. None means all tests in
+  that class are selected.
+
+  Args:
+    selected_tests: list of strings, list of tests to execute. E.g.
+
+    .. code-block:: python
+
+      ['FooTest', 'BarTest', 'BazTest.test_method_a', 'BazTest.test_method_b']
+
+  Returns:
+    A dict. Keys are test class names, values are lists of test names within
+    class. E.g. the example in `selected_tests` would translate to:
+
+    .. code-block:: python
+      {
+        'FooTest': None,
+        'BarTest': None,
+        'BazTest': ['test_method_a', 'test_method_b'],
+      }
+
+    This returns None if `selected_tests` is None.
+  """
+  if selected_tests is None:
+    return None
   test_class_name_to_tests = collections.OrderedDict()
   for test_name in selected_tests:
     if '.' in test_name:  # Has a test method
@@ -342,13 +388,4 @@
     else:  # No test method; run all tests in this class.
       test_class_name_to_tests[test_name] = None
 
-  # Now transform class names to class objects.
-  # Dict from test_name class name to instance.
-  class_name_to_class = {cls.__name__: cls for cls in test_classes}
-  for test_class_name, tests in test_class_name_to_tests.items():
-    test_class = class_name_to_class.get(test_class_name)
-    if not test_class:
-      raise Error('Unknown test_name class %s' % test_class_name)
-    class_to_tests[test_class] = tests
-
-  return class_to_tests
+  return test_class_name_to_tests
diff --git a/tests/mobly/base_suite_test.py b/tests/mobly/base_suite_test.py
new file mode 100644
index 0000000..43f9213
--- /dev/null
+++ b/tests/mobly/base_suite_test.py
@@ -0,0 +1,139 @@
+# Copyright 2024 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+from unittest import mock
+
+from mobly import base_suite
+from mobly import base_test
+from mobly import suite_runner
+from mobly import test_runner
+from mobly import config_parser
+from tests.lib import integration2_test
+from tests.lib import integration_test
+
+
+class FakeTest1(base_test.BaseTestClass):
+
+  def test_a(self):
+    pass
+
+  def test_b(self):
+    pass
+
+  def test_c(self):
+    pass
+
+
+class FakeTest2(base_test.BaseTestClass):
+
+  def test_2(self):
+    pass
+
+
+class FakeTestSuite(base_suite.BaseSuite):
+
+  def setup_suite(self, config):
+    self.add_test_class(FakeTest1, config)
+    self.add_test_class(FakeTest2, config)
+
+
+class FakeTestSuiteWithFilteredTests(base_suite.BaseSuite):
+
+  def setup_suite(self, config):
+    self.add_test_class(FakeTest1, config, ['test_a', 'test_b'])
+    self.add_test_class(FakeTest2, config, ['test_2'])
+
+
+class BaseSuiteTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.mock_config = mock.Mock(autospec=config_parser.TestRunConfig)
+    self.mock_test_runner = mock.Mock(autospec=test_runner.TestRunner)
+
+  def test_setup_suite(self):
+    suite = FakeTestSuite(self.mock_test_runner, self.mock_config)
+    suite.set_test_selector(None)
+
+    suite.setup_suite(self.mock_config)
+
+    self.mock_test_runner.add_test_class.assert_has_calls(
+        [
+            mock.call(self.mock_config, FakeTest1, mock.ANY, mock.ANY),
+            mock.call(self.mock_config, FakeTest2, mock.ANY, mock.ANY),
+        ],
+    )
+
+  def test_setup_suite_with_test_selector(self):
+    suite = FakeTestSuite(self.mock_test_runner, self.mock_config)
+    test_selector = {
+        'FakeTest1': ['test_a', 'test_b'],
+        'FakeTest2': None,
+    }
+
+    suite.set_test_selector(test_selector)
+    suite.setup_suite(self.mock_config)
+
+    self.mock_test_runner.add_test_class.assert_has_calls(
+        [
+            mock.call(
+                self.mock_config, FakeTest1, ['test_a', 'test_b'], mock.ANY
+            ),
+            mock.call(self.mock_config, FakeTest2, None, mock.ANY),
+        ],
+    )
+
+  def test_setup_suite_test_selector_takes_precedence(self):
+    suite = FakeTestSuiteWithFilteredTests(
+        self.mock_test_runner, self.mock_config
+    )
+    test_selector = {
+        'FakeTest1': ['test_a', 'test_c'],
+        'FakeTest2': None,
+    }
+
+    suite.set_test_selector(test_selector)
+    suite.setup_suite(self.mock_config)
+
+    self.mock_test_runner.add_test_class.assert_has_calls(
+        [
+            mock.call(
+                self.mock_config, FakeTest1, ['test_a', 'test_c'], mock.ANY
+            ),
+            mock.call(self.mock_config, FakeTest2, None, mock.ANY),
+        ],
+    )
+
+  def test_setup_suite_with_skip_test_class(self):
+    suite = FakeTestSuite(self.mock_test_runner, self.mock_config)
+    test_selector = {'FakeTest1': None}
+
+    suite.set_test_selector(test_selector)
+    suite.setup_suite(self.mock_config)
+
+    self.mock_test_runner.add_test_class.assert_has_calls(
+        [
+            mock.call(self.mock_config, FakeTest1, None, mock.ANY),
+        ],
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 08072d6..e4fad0a 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -23,6 +23,7 @@
 from mobly import base_suite
 from mobly import base_test
 from mobly import suite_runner
+from mobly import test_runner
 from tests.lib import integration2_test
 from tests.lib import integration_test
 
@@ -30,6 +31,9 @@
 class FakeTest1(base_test.BaseTestClass):
   pass
 
+  def test_a(self):
+    pass
+
 
 class SuiteRunnerTest(unittest.TestCase):
 
@@ -138,10 +142,16 @@
   @mock.patch('sys.exit')
   @mock.patch.object(suite_runner, '_find_suite_class', autospec=True)
   def test_run_suite_class(self, mock_find_suite_class, mock_exit):
+    tmp_file_path = self._gen_tmp_config_file()
+    mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
     mock_called = mock.MagicMock()
 
     class FakeTestSuite(base_suite.BaseSuite):
 
+      def set_test_selector(self, test_selector):
+        mock_called.set_test_selector(test_selector)
+        super().set_test_selector(test_selector)
+
       def setup_suite(self, config):
         mock_called.setup_suite()
         super().setup_suite(config)
@@ -153,20 +163,6 @@
 
     mock_find_suite_class.return_value = FakeTestSuite
 
-    tmp_file_path = os.path.join(self.tmp_dir, 'config.yml')
-    with io.open(tmp_file_path, 'w', encoding='utf-8') as f:
-      f.write(
-          """
-        TestBeds:
-          # A test bed where adb will find Android devices.
-          - Name: SampleTestBed
-            Controllers:
-              MagicDevice: '*'
-      """
-      )
-
-    mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
-
     with mock.patch.object(sys, 'argv', new=mock_cli_args):
       suite_runner.run_suite_class()
 
@@ -174,6 +170,75 @@
     mock_called.setup_suite.assert_called_once_with()
     mock_called.teardown_suite.assert_called_once_with()
     mock_exit.assert_not_called()
+    mock_called.set_test_selector.assert_called_once_with(None)
+
+  @mock.patch('sys.exit')
+  @mock.patch.object(suite_runner, '_find_suite_class', autospec=True)
+  @mock.patch.object(test_runner, 'TestRunner')
+  def test_run_suite_class_with_test_selection_by_class(
+      self, mock_test_runner_class, mock_find_suite_class, mock_exit
+  ):
+    mock_test_runner = mock_test_runner_class.return_value
+    mock_test_runner.results.is_all_pass = True
+    tmp_file_path = self._gen_tmp_config_file()
+    mock_cli_args = [
+        'test_binary',
+        f'--config={tmp_file_path}',
+        '--tests=FakeTest1',
+    ]
+    mock_called = mock.MagicMock()
+
+    class FakeTestSuite(base_suite.BaseSuite):
+
+      def set_test_selector(self, test_selector):
+        mock_called.set_test_selector(test_selector)
+        super().set_test_selector(test_selector)
+
+      def setup_suite(self, config):
+        self.add_test_class(FakeTest1)
+
+    mock_find_suite_class.return_value = FakeTestSuite
+
+    with mock.patch.object(sys, 'argv', new=mock_cli_args):
+      suite_runner.run_suite_class()
+
+    mock_called.set_test_selector.assert_called_once_with(
+        {'FakeTest1': None},
+    )
+
+  @mock.patch('sys.exit')
+  @mock.patch.object(suite_runner, '_find_suite_class', autospec=True)
+  @mock.patch.object(test_runner, 'TestRunner')
+  def test_run_suite_class_with_test_selection_by_method(
+      self, mock_test_runner_class, mock_find_suite_class, mock_exit
+  ):
+    mock_test_runner = mock_test_runner_class.return_value
+    mock_test_runner.results.is_all_pass = True
+    tmp_file_path = self._gen_tmp_config_file()
+    mock_cli_args = [
+        'test_binary',
+        f'--config={tmp_file_path}',
+        '--tests=FakeTest1.test_a',
+    ]
+    mock_called = mock.MagicMock()
+
+    class FakeTestSuite(base_suite.BaseSuite):
+
+      def set_test_selector(self, test_selector):
+        mock_called.set_test_selector(test_selector)
+        super().set_test_selector(test_selector)
+
+      def setup_suite(self, config):
+        self.add_test_class(FakeTest1)
+
+    mock_find_suite_class.return_value = FakeTestSuite
+
+    with mock.patch.object(sys, 'argv', new=mock_cli_args):
+      suite_runner.run_suite_class()
+
+    mock_called.set_test_selector.assert_called_once_with(
+        {'FakeTest1': ['test_a']},
+    )
 
   def test_print_test_names(self):
     mock_test_class = mock.MagicMock()
@@ -191,6 +256,20 @@
     mock_cls_instance._pre_run.side_effect = Exception('Something went wrong.')
     mock_cls_instance._clean_up.assert_called_once()
 
+  def _gen_tmp_config_file(self):
+    tmp_file_path = os.path.join(self.tmp_dir, 'config.yml')
+    with io.open(tmp_file_path, 'w', encoding='utf-8') as f:
+      f.write(
+          """
+        TestBeds:
+          # A test bed where adb will find Android devices.
+          - Name: SampleTestBed
+            Controllers:
+              MagicDevice: '*'
+      """
+      )
+    return tmp_file_path
+
 
 if __name__ == '__main__':
   unittest.main()