Add find_subclasses_in_module utility method. (#631)
diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index 962ef03..d32d304 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -146,12 +146,8 @@
Returns:
The test class in the test module.
"""
- test_classes = []
- main_module_members = sys.modules['__main__']
- for _, module_member in main_module_members.__dict__.items():
- if inspect.isclass(module_member):
- if issubclass(module_member, base_test.BaseTestClass):
- test_classes.append(module_member)
+ 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])
@@ -190,12 +186,10 @@
self.results: The test result object used to record the results of
this test run.
"""
-
class _TestRunInfo(object):
"""Identifies one test class to run, which tests to run, and config to
run it with.
"""
-
def __init__(self,
config,
test_class,
diff --git a/mobly/utils.py b/mobly/utils.py
index 17bcc88..a32c779 100644
--- a/mobly/utils.py
+++ b/mobly/utils.py
@@ -17,6 +17,7 @@
import base64
import concurrent.futures
import datetime
+import inspect
import io
import logging
import os
@@ -546,3 +547,23 @@
# Return directly if it's already a string.
return args
return ' '.join([pipes.quote(arg) for arg in args])
+
+
+def find_subclasses_in_module(base_classes, module):
+ """Finds the subclasses of the given classes in the given module.
+
+ Args:
+ base_classes: list of classes, the base classes to look for the
+ subclasses of in the module.
+ module: module, the module to look for the subclasses in.
+
+ Returns:
+ A list of all of the subclasses found in the module.
+ """
+ subclasses = []
+ for _, module_member in module.__dict__.items():
+ if inspect.isclass(module_member):
+ for base_class in base_classes:
+ if issubclass(module_member, base_class):
+ subclasses.append(module_member)
+ return subclasses
diff --git a/tests/lib/multiple_subclasses_module.py b/tests/lib/multiple_subclasses_module.py
new file mode 100755
index 0000000..0a82cfd
--- /dev/null
+++ b/tests/lib/multiple_subclasses_module.py
@@ -0,0 +1,32 @@
+# Copyright 2019 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.
+
+from mobly import base_test
+from mobly import test_runner
+
+class Subclass1Test(base_test.BaseTestClass):
+ pass
+
+class Subclass2Test(base_test.BaseTestClass):
+ pass
+
+class Subclass1Runner(test_runner.TestRunner):
+ pass
+
+class Subclass2Runner(test_runner.TestRunner):
+ pass
+
+class UnrelatedClass(object):
+ pass
+
diff --git a/tests/mobly/utils_test.py b/tests/mobly/utils_test.py
index 79dce6e..1c49377 100755
--- a/tests/mobly/utils_test.py
+++ b/tests/mobly/utils_test.py
@@ -25,7 +25,16 @@
import portpicker
import psutil
+from mobly import base_instrumentation_test
+
+from mobly import base_test
+from mobly import signals
+from mobly import test_runner
from mobly import utils
+from tests.lib import integration_test
+from tests.lib import mock_controller
+from tests.lib import mock_instrumentation_test
+from tests.lib import multiple_subclasses_module
MOCK_AVAILABLE_PORT = 5
@@ -34,7 +43,6 @@
"""This test class has unit tests for the implementation of everything
under mobly.utils.
"""
-
def setUp(self):
system = platform.system()
self.tmp_dir = tempfile.mkdtemp()
@@ -92,13 +100,12 @@
mock_proc = mock_Popen.return_value
mock_proc.communicate.return_value = ('fake_out', 'fake_err')
mock_proc.returncode = 127
- out = utils.run_command(
- mock_command,
- stdout=mock_stdout,
- stderr=mock_stderr,
- shell=mock_shell,
- timeout=mock_timeout,
- env=mock_env)
+ out = utils.run_command(mock_command,
+ stdout=mock_stdout,
+ stderr=mock_stderr,
+ shell=mock_shell,
+ timeout=mock_timeout,
+ env=mock_env)
self.assertEqual(out, (127, 'fake_out', 'fake_err'))
mock_Popen.assert_called_with(
mock_command,
@@ -215,27 +222,24 @@
expected_base64_encoding = u'SGVsbG93IHdvcmxkIQ=='
with io.open(tmp_file_path, 'wb') as f:
f.write(b'Hellow world!')
- self.assertEqual(
- utils.load_file_to_base64_str(tmp_file_path),
- expected_base64_encoding)
+ self.assertEqual(utils.load_file_to_base64_str(tmp_file_path),
+ expected_base64_encoding)
def test_load_file_to_base64_str_reads_text_file_as_base64_string(self):
tmp_file_path = os.path.join(self.tmp_dir, 'b64.bin')
expected_base64_encoding = u'SGVsbG93IHdvcmxkIQ=='
with io.open(tmp_file_path, 'w', encoding='utf-8') as f:
f.write(u'Hellow world!')
- self.assertEqual(
- utils.load_file_to_base64_str(tmp_file_path),
- expected_base64_encoding)
+ self.assertEqual(utils.load_file_to_base64_str(tmp_file_path),
+ expected_base64_encoding)
def test_load_file_to_base64_str_reads_unicode_file_as_base64_string(self):
tmp_file_path = os.path.join(self.tmp_dir, 'b64.bin')
expected_base64_encoding = u'6YCa'
with io.open(tmp_file_path, 'w', encoding='utf-8') as f:
f.write(u'\u901a')
- self.assertEqual(
- utils.load_file_to_base64_str(tmp_file_path),
- expected_base64_encoding)
+ self.assertEqual(utils.load_file_to_base64_str(tmp_file_path),
+ expected_base64_encoding)
def test_cli_cmd_to_string(self):
cmd = ['"adb"', 'a b', 'c//']
@@ -243,6 +247,50 @@
cmd = 'adb -s meme do something ab_cd'
self.assertEqual(utils.cli_cmd_to_string(cmd), cmd)
+ def test_find_subclasses_in_module_when_one_subclass(self):
+ subclasses = utils.find_subclasses_in_module([base_test.BaseTestClass],
+ integration_test)
+ self.assertEqual(len(subclasses), 1)
+ self.assertEqual(subclasses[0], integration_test.IntegrationTest)
+
+ def test_find_subclasses_in_module_when_indirect_subclass(self):
+ subclasses = utils.find_subclasses_in_module([base_test.BaseTestClass],
+ mock_instrumentation_test)
+ self.assertEqual(len(subclasses), 1)
+ self.assertEqual(subclasses[0],
+ mock_instrumentation_test.MockInstrumentationTest)
+
+ def test_find_subclasses_in_module_when_no_subclasses(self):
+ subclasses = utils.find_subclasses_in_module([base_test.BaseTestClass],
+ mock_controller)
+ self.assertEqual(len(subclasses), 0)
+
+ def test_find_subclasses_in_module_when_multiple_subclasses(self):
+ subclasses = utils.find_subclasses_in_module(
+ [base_test.BaseTestClass], multiple_subclasses_module)
+ self.assertEqual(len(subclasses), 2)
+ self.assertIn(multiple_subclasses_module.Subclass1Test, subclasses)
+ self.assertIn(multiple_subclasses_module.Subclass2Test, subclasses)
+
+ def test_find_subclasses_in_module_when_multiple_base_classes(self):
+ subclasses = utils.find_subclasses_in_module(
+ [base_test.BaseTestClass, test_runner.TestRunner],
+ multiple_subclasses_module)
+ self.assertEqual(len(subclasses), 4)
+ self.assertIn(multiple_subclasses_module.Subclass1Test, subclasses)
+ self.assertIn(multiple_subclasses_module.Subclass2Test, subclasses)
+ self.assertIn(multiple_subclasses_module.Subclass1Runner, subclasses)
+ self.assertIn(multiple_subclasses_module.Subclass2Runner, subclasses)
+
+ def test_find_subclasses_in_module_when_only_some_base_classes_present(
+ self):
+ subclasses = utils.find_subclasses_in_module(
+ [signals.TestSignal, test_runner.TestRunner],
+ multiple_subclasses_module)
+ self.assertEqual(len(subclasses), 2)
+ self.assertIn(multiple_subclasses_module.Subclass1Runner, subclasses)
+ self.assertIn(multiple_subclasses_module.Subclass2Runner, subclasses)
+
if __name__ == '__main__':
unittest.main()
