From f596d09f2945ded6b284c995de1dcbb0c04a3923 Mon Sep 17 00:00:00 2001
From: winterfroststrom <winterfroststrom@users.noreply.github.com>
Date: Mon, 23 Sep 2019 14:00:23 -0700
Subject: [PATCH] Add find_subclass_in_module utility. (#632)

---
 mobly/test_runner.py            | 15 +++++++++------
 mobly/utils.py                  | 21 +++++++++++++++++++++
 tests/mobly/test_runner_test.py | 17 +++++++++++++++++
 tests/mobly/utils_test.py       | 28 ++++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index d32d3041..4dabfb24 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -145,14 +145,17 @@ def _find_test_class():
 
     Returns:
         The test class in the test module.
+
+    Raises:
+      SystemExit: Raised if the number of test classes is not exactly one.
     """
-    test_classes = utils.find_subclasses_in_module(base_test.BaseTestClass,
-                                                   sys.modules['__main__'])
-    if len(test_classes) != 1:
-        logging.error('Expected 1 test class per file, found %s.',
-                      [t.__name__ for t in test_classes])
+    try:
+        return utils.find_subclass_in_module(base_test.BaseTestClass,
+                                             sys.modules['__main__'])
+    except ValueError:
+        logging.exception('Exactly one subclass of `base_test.BaseTestClass`'
+                          ' should be in the main file.')
         sys.exit(1)
-    return test_classes[0]
 
 
 def _print_test_names(test_class):
diff --git a/mobly/utils.py b/mobly/utils.py
index a32c7799..2a60bae1 100644
--- a/mobly/utils.py
+++ b/mobly/utils.py
@@ -567,3 +567,24 @@ def find_subclasses_in_module(base_classes, module):
                 if issubclass(module_member, base_class):
                     subclasses.append(module_member)
     return subclasses
+
+
+def find_subclass_in_module(base_class, module):
+    """Finds the single subclass of the given base class in the given module.
+
+    Args:
+      base_class: class, the base class to look for a subclass of in the module.
+      module: module, the module to look for the single subclass in.
+
+    Returns:
+      The single subclass of the given base class.
+
+    Raises:
+      ValueError: If the number of subclasses found was not exactly one.
+    """
+    subclasses = find_subclasses_in_module([base_class], module)
+    if len(subclasses) != 1:
+        raise ValueError('Expected 1 subclass of %s per module, found %s.' %
+                         (base_class.__name__,
+                          [subclass.__name__ for subclass in subclasses]))
+    return subclasses[0]
diff --git a/tests/mobly/test_runner_test.py b/tests/mobly/test_runner_test.py
index 90e92681..4074891b 100755
--- a/tests/mobly/test_runner_test.py
+++ b/tests/mobly/test_runner_test.py
@@ -32,6 +32,7 @@
 from tests.lib import integration_test
 from tests.lib import integration2_test
 from tests.lib import integration3_test
+from tests.lib import multiple_subclasses_module
 
 
 class TestRunnerTest(unittest.TestCase):
@@ -342,6 +343,22 @@ def test_main_with_failures(self, mock_exit, mock_find_test):
         test_runner.main(['-c', tmp_file_path])
         mock_exit.assert_called_once_with(1)
 
+    def test__find_test_class_when_one_test_class(self):
+        with mock.patch.dict('sys.modules', __main__=integration_test):
+            test_class = test_runner._find_test_class()
+            self.assertEqual(test_class, integration_test.IntegrationTest)
+
+    def test__find_test_class_when_no_test_class(self):
+        with self.assertRaises(SystemExit):
+            with mock.patch.dict('sys.modules', __main__=mock_controller):
+                test_class = test_runner._find_test_class()
+
+    def test__find_test_class_when_multiple_test_classes(self):
+        with self.assertRaises(SystemExit):
+            with mock.patch.dict('sys.modules',
+                                 __main__=multiple_subclasses_module):
+                test_class = test_runner._find_test_class()
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/mobly/utils_test.py b/tests/mobly/utils_test.py
index 1c493778..714c2386 100755
--- a/tests/mobly/utils_test.py
+++ b/tests/mobly/utils_test.py
@@ -291,6 +291,34 @@ def test_find_subclasses_in_module_when_only_some_base_classes_present(
         self.assertIn(multiple_subclasses_module.Subclass1Runner, subclasses)
         self.assertIn(multiple_subclasses_module.Subclass2Runner, subclasses)
 
+    def test_find_subclass_in_module_when_one_subclass(self):
+        subclass = utils.find_subclass_in_module(base_test.BaseTestClass,
+                                                 integration_test)
+        self.assertEqual(subclass, integration_test.IntegrationTest)
+
+    def test_find_subclass_in_module_when_indirect_subclass(self):
+        subclass = utils.find_subclass_in_module(base_test.BaseTestClass,
+                                                 mock_instrumentation_test)
+        self.assertEqual(subclass,
+                         mock_instrumentation_test.MockInstrumentationTest)
+
+    def test_find_subclass_in_module_when_no_subclasses(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                '.*Expected 1 subclass of BaseTestClass per module, found'
+                r' \[\].*'):
+            _ = utils.find_subclass_in_module(base_test.BaseTestClass,
+                                              mock_controller)
+
+    def test_find_subclass_in_module_when_multiple_subclasses(self):
+        with self.assertRaisesRegex(
+                ValueError,
+                '.*Expected 1 subclass of BaseTestClass per module, found'
+                r' \[(\'Subclass1Test\', \'Subclass2Test\''
+                r'|\'Subclass2Test\', \'Subclass1Test\')\].*'):
+            _ = utils.find_subclass_in_module(base_test.BaseTestClass,
+                                              multiple_subclasses_module)
+
 
 if __name__ == '__main__':
     unittest.main()