Repository: google/mobly

Issue:
Move `register_controller` into `base_test`
For historic reasons, `register_controller` lives under `test_runner`.
Back then the consideration was, one day we may want to let controller objects persist across multiple test classes.

This causes multiple organizational problems and incurs complexity.

In hindsight that is a very bad idea and controller object lifecycles should be contained within each class. So it is time to push controller management into `base_test`.

Discussion:
xpconanfan: CI failure it not reproducible locally in py2.7, 3.4, 3.5, 3.6
I don't know what Travis is doing...
dthkao: https://stackoverflow.com/questions/24802740/unnamed-python-objects-have-the-same-id

Looks like the test was based on a bad assumption to begin wth. The two objects you're id()-ing have non-overlapping scope so no guarantee of unique mem address.

Try making them attributes of the outer test (although this relies on the mock controller destroy not actually doing anything)?
winterfroststrom: I re-ran the build, so it's probably a test flake. Probably per the reason dthkao@ mentioned.
Is there an easy way to rewrite the test to avoid this?

Patch:
diff --git a/mobly/config_parser.py b/mobly/config_parser.py
index b873af0..2278ca6 100644
--- a/mobly/config_parser.py
+++ b/mobly/config_parser.py
@@ -166,8 +166,6 @@ class TestRunConfig(object):
         controller_configs: dict, configs used for instantiating controller
             objects.
         user_params: dict, all the parameters to be consumed by the test logic.
-        register_controller: func, used by test classes to register controller
-            modules.
         summary_writer: records.TestSummaryWriter, used to write elements to
             the test result summary file.
         test_class_name_suffix: string, suffix to append to the class name for
@@ -180,7 +178,6 @@ class TestRunConfig(object):
         self.test_bed_name = None
         self.controller_configs = None
         self.user_params = None
-        self.register_controller = None
         self.summary_writer = None
         self.test_class_name_suffix = None
 
@@ -192,5 +189,4 @@ class TestRunConfig(object):
     def __str__(self):
         content = dict(self.__dict__)
         content.pop('summary_writer')
-        content.pop('register_controller')
         return pprint.pformat(content)
diff --git a/mobly/records.py b/mobly/records.py
index ad37cad..051ff9c 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -488,10 +488,10 @@ class TestResult(object):
                 setattr(sum_result, name, l_value + r_value)
             elif isinstance(r_value, dict):
                 # '+' operator for TestResult is only valid when multiple
-                # TestResult objs were created in the same test run, which means
-                # the controller info would be the same across all of them.
+                # TestResult objs were created in the same test run, use the
+                # r-value which is more up to date.
                 # TODO(xpconanfan): have a better way to validate this situation.
-                setattr(sum_result, name, l_value)
+                setattr(sum_result, name, r_value)
         return sum_result
 
     def add_record(self, record):


Test patch:
diff --git a/mobly/base_test.py b/mobly/base_test.py
index e64c437..e4139d9 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -46,6 +46,30 @@ class Error(Exception):
     """Raised for exceptions that occured in BaseTestClass."""
 
 
+def _verify_controller_module(module):
+    """Verifies a module object follows the required interface for
+    controllers.
+
+    Args:
+        module: An object that is a controller module. This is usually
+            imported with import statements or loaded by importlib.
+
+    Raises:
+        ControllerError: if the module does not match the Mobly controller
+            interface, or one of the required members is null.
+    """
+    required_attributes = ('create', 'destroy', 'MOBLY_CONTROLLER_CONFIG_NAME')
+    for attr in required_attributes:
+        if not hasattr(module, attr):
+            raise signals.ControllerError(
+                'Module %s missing required controller module attribute'
+                ' %s.' % (module.__name__, attr))
+        if not getattr(module, attr):
+            raise signals.ControllerError(
+                'Controller interface %s in %s cannot be null.' %
+                (attr, module.__name__))
+
+
 class BaseTestClass(object):
     """Base class for all test classes to inherit from.
 
@@ -77,8 +101,6 @@ class BaseTestClass(object):
             objects.
         user_params: dict, custom parameters from user, to be consumed by
             the test logic.
-        register_controller: func, used by test classes to register
-            controller modules.
     """
 
     TAG = None
@@ -106,12 +128,14 @@ class BaseTestClass(object):
         self.controller_configs = configs.controller_configs
         self.test_bed_name = configs.test_bed_name
         self.user_params = configs.user_params
-        self.register_controller = configs.register_controller
         self.results = records.TestResult()
         self.summary_writer = configs.summary_writer
         # Deprecated, use `self.current_test_info.name`.
         self.current_test_name = None
         self._generated_test_table = collections.OrderedDict()
+        # Controller object management.
+        self._controller_registry = {}
+        self._controller_destructors = {}
 
     def __enter__(self):
         return self
@@ -171,6 +195,157 @@ class BaseTestClass(object):
                 logging.warning('Missing optional user param "%s" in '
                                 'configuration, continue.', name)
 
+    def register_controller(self, module, required=True, min_number=1):
+        """Loads a controller module and returns its loaded devices.
+
+        A Mobly controller module is a Python lib that can be used to control
+        a device, service, or equipment. To be Mobly compatible, a controller
+        module needs to have the following members:
+
+        ```
+        def create(configs):
+            [Required] Creates controller objects from configurations.
+
+            Args:
+                configs: A list of serialized data like string/dict. Each
+                    element of the list is a configuration for a controller
+                    object.
+
+            Returns:
+                A list of objects.
+
+        def destroy(objects):
+            [Required] Destroys controller objects created by the create
+            function. Each controller object shall be properly cleaned up
+            and all the resources held should be released, e.g. memory
+            allocation, sockets, file handlers etc.
+
+            Args:
+                A list of controller objects created by the create function.
+
+        def get_info(objects):
+            [Optional] Gets info from the controller objects used in a test
+            run. The info will be included in test_summary.yaml under
+            the key 'ControllerInfo'. Such information could include unique
+            ID, version, or anything that could be useful for describing the
+            test bed and debugging.
+
+            Args:
+                objects: A list of controller objects created by the create
+                    function.
+
+            Returns:
+                A list of json serializable objects, each represents the
+                info of a controller object. The order of the info object
+                should follow that of the input objects.
+        ```
+
+        Registering a controller module declares a test class's dependency the
+        controller. If the module config exists and the module matches the
+        controller interface, controller objects will be instantiated with
+        corresponding configs. The module should be imported first.
+
+        Args:
+            module: A module that follows the controller module interface.
+            required: A bool. If True, failing to register the specified
+                controller module raises exceptions. If False, the objects
+                failed to instantiate will be skipped.
+            min_number: An integer that is the minimum number of controller
+                objects to be created. Default is one, since you should not
+                register a controller module without expecting at least one
+                object.
+
+        Returns:
+            A list of controller objects instantiated from controller_module, or
+            None if no config existed for this controller and it was not a
+            required controller.
+
+        Raises:
+            ControllerError:
+                * The controller module has already been registered.
+                * The actual number of objects instantiated is less than the
+                * `min_number`.
+                * `required` is True and no corresponding config can be found.
+                * Any other error occurred in the registration process.
+        """
+        _verify_controller_module(module)
+        # Use the module's name as the ref name
+        module_ref_name = module.__name__.split('.')[-1]
+        if module_ref_name in self._controller_registry:
+            raise signals.ControllerError(
+                'Controller module %s has already been registered. It cannot '
+                'be registered again.' % module_ref_name)
+        # Create controller objects.
+        create = module.create
+        module_config_name = module.MOBLY_CONTROLLER_CONFIG_NAME
+        if module_config_name not in self.controller_configs:
+            if required:
+                raise signals.ControllerError(
+                    'No corresponding config found for %s' %
+                    module_config_name)
+            logging.warning(
+                'No corresponding config found for optional controller %s',
+                module_config_name)
+            return None
+        try:
+            # Make a deep copy of the config to pass to the controller module,
+            # in case the controller module modifies the config internally.
+            original_config = self.controller_configs[module_config_name]
+            controller_config = copy.deepcopy(original_config)
+            objects = create(controller_config)
+        except:
+            logging.exception(
+                'Failed to initialize objects for controller %s, abort!',
+                module_config_name)
+            raise
+        if not isinstance(objects, list):
+            raise signals.ControllerError(
+                'Controller module %s did not return a list of objects, abort.'
+                % module_ref_name)
+        # Check we got enough controller objects to continue.
+        actual_number = len(objects)
+        if actual_number < min_number:
+            module.destroy(objects)
+            raise signals.ControllerError(
+                'Expected to get at least %d controller objects, got %d.' %
+                (min_number, actual_number))
+        # Save a shallow copy of the list for internal usage, so tests can't
+        # affect internal registry by manipulating the object list.
+        self._controller_registry[module_ref_name] = copy.copy(objects)
+        # Collect controller information and write to test result.
+        # Implementation of 'get_info' is optional for a controller module.
+        if hasattr(module, 'get_info'):
+            controller_info = module.get_info(copy.copy(objects))
+            logging.debug('Controller %s: %s', module_config_name,
+                          controller_info)
+            self.results.add_controller_info(module_config_name,
+                                             controller_info)
+        else:
+            logging.warning('No optional debug info found for controller %s. '
+                            'To provide it, implement get_info in this '
+                            'controller module.', module_config_name)
+        logging.debug('Found %d objects for controller %s', len(objects),
+                      module_config_name)
+        destroy_func = module.destroy
+        self._controller_destructors[module_ref_name] = destroy_func
+        return objects
+
+    def _unregister_controllers(self):
+        """Destroy controller objects and clear internal registry.
+
+        This will be called after each test class.
+        """
+        # TODO(xpconanfan): actually record these errors instead of just
+        # logging them.
+        for name, destroy in self._controller_destructors.items():
+            try:
+                logging.debug('Destroying %s.', name)
+                destroy(self._controller_registry[name])
+            except:
+                logging.exception('Exception occurred destroying %s.', name)
+        self._controller_registry = {}
+        self._controller_destructors = {}
+
     def _setup_generated_tests(self):
         """Proxy function to guarantee the base implementation of
         setup_generated_tests is called.
@@ -727,6 +902,7 @@ class BaseTestClass(object):
             raise e
         finally:
             self._teardown_class()
+            self._unregister_controllers()
             logging.info('Summary for test class %s: %s', self.TAG,
                          self.results.summary_str())
 
diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index e4fa9d8..9c57715 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -17,8 +17,6 @@ from future import standard_library
 standard_library.install_aliases()
 
 import argparse
-import copy
-import functools
 import inspect
 import logging
 import os
@@ -180,73 +178,6 @@ def _print_test_names(test_class):
         print(name)
 
 
-def verify_controller_module(module):
-    """Verifies a module object follows the required interface for
-    controllers.
-
-    A Mobly controller module is a Python lib that can be used to control
-    a device, service, or equipment. To be Mobly compatible, a controller
-    module needs to have the following members:
-
-        def create(configs):
-            [Required] Creates controller objects from configurations.
-
-            Args:
-                configs: A list of serialized data like string/dict. Each
-                    element of the list is a configuration for a controller
-                    object.
-            Returns:
-                A list of objects.
-
-        def destroy(objects):
-            [Required] Destroys controller objects created by the create
-            function. Each controller object shall be properly cleaned up
-            and all the resources held should be released, e.g. memory
-            allocation, sockets, file handlers etc.
-
-            Args:
-                A list of controller objects created by the create function.
-
-        def get_info(objects):
-            [Optional] Gets info from the controller objects used in a test
-            run. The info will be included in test_summary.yaml under
-            the key 'ControllerInfo'. Such information could include unique
-            ID, version, or anything that could be useful for describing the
-            test bed and debugging.
-
-            Args:
-                objects: A list of controller objects created by the create
-                    function.
-            Returns:
-                A list of json serializable objects, each represents the
-                info of a controller object. The order of the info object
-                should follow that of the input objects.
-
-    Registering a controller module declares a test class's dependency the
-    controller. If the module config exists and the module matches the
-    controller interface, controller objects will be instantiated with
-    corresponding configs. The module should be imported first.
-
-    Args:
-        module: An object that is a controller module. This is usually
-            imported with import statements or loaded by importlib.
-
-    Raises:
-        ControllerError: if the module does not match the Mobly controller
-            interface, or one of the required members is null.
-    """
-    required_attributes = ('create', 'destroy', 'MOBLY_CONTROLLER_CONFIG_NAME')
-    for attr in required_attributes:
-        if not hasattr(module, attr):
-            raise signals.ControllerError(
-                'Module %s missing required controller module attribute'
-                ' %s.' % (module.__name__, attr))
-        if not getattr(module, attr):
-            raise signals.ControllerError(
-                'Controller interface %s in %s cannot be null.' %
-                (attr, module.__name__))
-
-
 class TestRunner(object):
     """The class that instantiates test classes, executes tests, and
     report results.
@@ -288,10 +219,6 @@ class TestRunner(object):
         self.results = records.TestResult()
         self._test_run_infos = []
 
-        # Controller management. These members will be updated for each class.
-        self._controller_registry = {}
-        self._controller_destructors = {}
-
         self._log_path = None
 
     def setup_logger(self):
@@ -412,8 +339,6 @@ class TestRunner(object):
                 # Set up the test-specific config
                 test_config = test_run_info.config.copy()
                 test_config.log_path = self._log_path
-                test_config.register_controller = functools.partial(
-                    self._register_controller, test_config)
                 test_config.summary_writer = summary_writer
                 test_config.test_class_name_suffix = test_run_info.test_class_name_suffix
                 try:
@@ -425,8 +350,6 @@ class TestRunner(object):
                     logging.warning(
                         'Abort all subsequent test classes. Reason: %s', e)
                     raise
-                finally:
-                    self._unregister_controllers()
         finally:
             # Write controller info and summary to summary file.
             summary_writer.dump(self.results.controller_info,
@@ -439,112 +362,3 @@ class TestRunner(object):
                 self.results.summary_str())
             logging.info(msg.strip())
             self._teardown_logger()
-
-    def _register_controller(self, config, module, required=True,
-                             min_number=1):
-        """Loads a controller module and returns its loaded devices.
-
-        See the docstring of verify_controller_module() for a description of
-        what API a controller module must implement to be compatible with this
-        method.
-
-        Args:
-            config: A config_parser.TestRunConfig object.
-            module: A module that follows the controller module interface.
-            required: A bool. If True, failing to register the specified
-                controller module raises exceptions. If False, the objects
-                failed to instantiate will be skipped.
-            min_number: An integer that is the minimum number of controller
-                objects to be created. Default is one, since you should not
-                register a controller module without expecting at least one
-                object.
-
-        Returns:
-            A list of controller objects instantiated from controller_module, or
-            None if no config existed for this controller and it was not a
-            required controller.
-
-        Raises:
-            ControllerError:
-                * The controller module has already been registered.
-                * The actual number of objects instantiated is less than the
-                * `min_number`.
-                * `required` is True and no corresponding config can be found.
-                * Any other error occurred in the registration process.
-
-        """
-        verify_controller_module(module)
-        # Use the module's name as the ref name
-        module_ref_name = module.__name__.split('.')[-1]
-        if module_ref_name in self._controller_registry:
-            raise signals.ControllerError(
-                'Controller module %s has already been registered. It cannot '
-                'be registered again.' % module_ref_name)
-        # Create controller objects.
-        create = module.create
-        module_config_name = module.MOBLY_CONTROLLER_CONFIG_NAME
-        if module_config_name not in config.controller_configs:
-            if required:
-                raise signals.ControllerError(
-                    'No corresponding config found for %s' %
-                    module_config_name)
-            logging.warning(
-                'No corresponding config found for optional controller %s',
-                module_config_name)
-            return None
-        try:
-            # Make a deep copy of the config to pass to the controller module,
-            # in case the controller module modifies the config internally.
-            original_config = config.controller_configs[module_config_name]
-            controller_config = copy.deepcopy(original_config)
-            objects = create(controller_config)
-        except:
-            logging.exception(
-                'Failed to initialize objects for controller %s, abort!',
-                module_config_name)
-            raise
-        if not isinstance(objects, list):
-            raise signals.ControllerError(
-                'Controller module %s did not return a list of objects, abort.'
-                % module_ref_name)
-        # Check we got enough controller objects to continue.
-        actual_number = len(objects)
-        if actual_number < min_number:
-            module.destroy(objects)
-            raise signals.ControllerError(
-                'Expected to get at least %d controller objects, got %d.' %
-                (min_number, actual_number))
-        # Save a shallow copy of the list for internal usage, so tests can't
-        # affect internal registry by manipulating the object list.
-        self._controller_registry[module_ref_name] = copy.copy(objects)
-        # Collect controller information and write to test result.
-        # Implementation of 'get_info' is optional for a controller module.
-        if hasattr(module, 'get_info'):
-            controller_info = module.get_info(copy.copy(objects))
-            logging.debug('Controller %s: %s', module_config_name,
-                          controller_info)
-            self.results.add_controller_info(module_config_name,
-                                             controller_info)
-        else:
-            logging.warning('No optional debug info found for controller %s. '
-                            'To provide it, implement get_info in this '
-                            'controller module.', module_config_name)
-        logging.debug('Found %d objects for controller %s', len(objects),
-                      module_config_name)
-        destroy_func = module.destroy
-        self._controller_destructors[module_ref_name] = destroy_func
-        return objects
-
-    def _unregister_controllers(self):
-        """Destroy controller objects and clear internal registry.
-
-        This will be called after each test class.
-        """
-        for name, destroy in self._controller_destructors.items():
-            try:
-                logging.debug('Destroying %s.', name)
-                destroy(self._controller_registry[name])
-            except:
-                logging.exception('Exception occurred destroying %s.', name)
-        self._controller_registry = {}
-        self._controller_destructors = {}
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index d029219..898ab5d 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -29,6 +29,7 @@ from mobly import records
 from mobly import signals
 
 from tests.lib import utils
+from tests.lib import mock_controller
 
 MSG_EXPECTED_EXCEPTION = "This is an expected exception."
 MSG_EXPECTED_TEST_FAILURE = "This is an expected test failure."
@@ -45,6 +46,15 @@ class SomeError(Exception):
     """A custom exception class used for tests in this module."""
 
 
+class MockEmptyBaseTest(base_test.BaseTestClass):
+    """Stub used to test functionalities not specific to a class
+    implementation.
+    """
+
+    def test_func(self):
+        pass
+
+
 class BaseTestTest(unittest.TestCase):
     def setUp(self):
         self.tmp_dir = tempfile.mkdtemp()
@@ -52,6 +62,7 @@ class BaseTestTest(unittest.TestCase):
         self.summary_file = os.path.join(self.tmp_dir, 'summary.yaml')
         self.mock_test_cls_configs.summary_writer = records.TestSummaryWriter(
             self.summary_file)
+        self.mock_test_cls_configs.controller_configs = {}
         self.mock_test_cls_configs.log_path = self.tmp_dir
         self.mock_test_cls_configs.user_params = {"some_param": "hahaha"}
         self.mock_test_cls_configs.reporter = mock.MagicMock()
@@ -1773,6 +1784,113 @@ class BaseTestTest(unittest.TestCase):
                 self.assertIsNotNone(c['timestamp'])
         self.assertTrue(hit)
 
+    def test_register_controller_no_config(self):
+        bt_cls = MockEmptyBaseTest(self.mock_test_cls_configs)
+        with self.assertRaisesRegex(signals.ControllerError,
+                                    'No corresponding config found for'):
+            bt_cls.register_controller(mock_controller)
+
+    def test_register_controller_no_config_for_not_required(self):
+        bt_cls = MockEmptyBaseTest(self.mock_test_cls_configs)
+        self.assertIsNone(
+            bt_cls.register_controller(mock_controller, required=False))
+
+    def test_register_controller_dup_register(self):
+        """Verifies correctness of registration, internal tally of controllers
+        objects, and the right error happen when a controller module is
+        registered twice.
+        """
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_test_config.controller_configs = {
+            mock_ctrlr_config_name: ['magic1', 'magic2']
+        }
+        bt_cls = MockEmptyBaseTest(mock_test_config)
+        bt_cls.register_controller(mock_controller)
+        registered_name = 'mock_controller'
+        self.assertTrue(registered_name in bt_cls._controller_registry)
+        mock_ctrlrs = bt_cls._controller_registry[registered_name]
+        self.assertEqual(mock_ctrlrs[0].magic, 'magic1')
+        self.assertEqual(mock_ctrlrs[1].magic, 'magic2')
+        self.assertTrue(bt_cls._controller_destructors[registered_name])
+        expected_msg = 'Controller module .* has already been registered.'
+        with self.assertRaisesRegex(signals.ControllerError, expected_msg):
+            bt_cls.register_controller(mock_controller)
+
+    def test_register_controller_no_get_info(self):
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        get_info = getattr(mock_controller, 'get_info')
+        delattr(mock_controller, 'get_info')
+        try:
+            mock_test_config.controller_configs = {
+                mock_ctrlr_config_name: ['magic1', 'magic2']
+            }
+            bt_cls = MockEmptyBaseTest(mock_test_config)
+            bt_cls.register_controller(mock_controller)
+            self.assertEqual(bt_cls.results.controller_info, {})
+        finally:
+            setattr(mock_controller, 'get_info', get_info)
+
+    def test_register_controller_return_value(self):
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_test_config.controller_configs = {
+            mock_ctrlr_config_name: ['magic1', 'magic2']
+        }
+        bt_cls = MockEmptyBaseTest(mock_test_config)
+        magic_devices = bt_cls.register_controller(mock_controller)
+        self.assertEqual(magic_devices[0].magic, 'magic1')
+        self.assertEqual(magic_devices[1].magic, 'magic2')
+
+    def test_register_controller_change_return_value(self):
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_test_config.controller_configs = {
+            mock_ctrlr_config_name: ['magic1', 'magic2']
+        }
+        bt_cls = MockEmptyBaseTest(mock_test_config)
+        magic_devices = bt_cls.register_controller(mock_controller)
+        magic1 = magic_devices.pop(0)
+        self.assertIs(magic1,
+                      bt_cls._controller_registry['mock_controller'][0])
+        self.assertEqual(
+            len(bt_cls._controller_registry['mock_controller']), 2)
+
+    def test_register_controller_less_than_min_number(self):
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_test_config.controller_configs = {
+            mock_ctrlr_config_name: ['magic1', 'magic2']
+        }
+        bt_cls = MockEmptyBaseTest(mock_test_config)
+        expected_msg = 'Expected to get at least 3 controller objects, got 2.'
+        with self.assertRaisesRegex(signals.ControllerError, expected_msg):
+            bt_cls.register_controller(mock_controller, min_number=3)
+
+    def test_verify_controller_module(self):
+        base_test._verify_controller_module(mock_controller)
+
+    def test_verify_controller_module_null_attr(self):
+        try:
+            tmp = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+            mock_controller.MOBLY_CONTROLLER_CONFIG_NAME = None
+            msg = 'Controller interface .* in .* cannot be null.'
+            with self.assertRaisesRegex(signals.ControllerError, msg):
+                base_test._verify_controller_module(mock_controller)
+        finally:
+            mock_controller.MOBLY_CONTROLLER_CONFIG_NAME = tmp
+
+    def test_verify_controller_module_missing_attr(self):
+        try:
+            tmp = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+            delattr(mock_controller, 'MOBLY_CONTROLLER_CONFIG_NAME')
+            msg = 'Module .* missing required controller module attribute'
+            with self.assertRaisesRegex(signals.ControllerError, msg):
+                base_test._verify_controller_module(mock_controller)
+        finally:
+            setattr(mock_controller, 'MOBLY_CONTROLLER_CONFIG_NAME', tmp)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/mobly/test_runner_test.py b/tests/mobly/test_runner_test.py
index b369231..d86f60c 100755
--- a/tests/mobly/test_runner_test.py
+++ b/tests/mobly/test_runner_test.py
@@ -16,7 +16,6 @@ import io
 import logging
 import mock
 import os
-import platform
 import re
 import shutil
 import tempfile
@@ -56,93 +55,6 @@ class TestRunnerTest(unittest.TestCase):
     def tearDown(self):
         shutil.rmtree(self.tmp_dir)
 
-    def test_register_controller_no_config(self):
-        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-        with self.assertRaisesRegex(signals.ControllerError,
-                                    'No corresponding config found for'):
-            tr._register_controller(self.base_mock_test_config,
-                                    mock_controller)
-
-    def test_register_controller_no_config_no_register(self):
-        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-        self.assertIsNone(
-            tr._register_controller(
-                self.base_mock_test_config, mock_controller, required=False))
-
-    def test_register_controller_dup_register(self):
-        """Verifies correctness of registration, internal tally of controllers
-        objects, and the right error happen when a controller module is
-        registered twice.
-        """
-        mock_test_config = self.base_mock_test_config.copy()
-        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-        mock_test_config.controller_configs = {
-            mock_ctrlr_config_name: ['magic1', 'magic2']
-        }
-        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-        tr._register_controller(mock_test_config, mock_controller)
-        registered_name = 'mock_controller'
-        self.assertTrue(registered_name in tr._controller_registry)
-        mock_ctrlrs = tr._controller_registry[registered_name]
-        self.assertEqual(mock_ctrlrs[0].magic, 'magic1')
-        self.assertEqual(mock_ctrlrs[1].magic, 'magic2')
-        self.assertTrue(tr._controller_destructors[registered_name])
-        expected_msg = 'Controller module .* has already been registered.'
-        with self.assertRaisesRegex(signals.ControllerError, expected_msg):
-            tr._register_controller(mock_test_config, mock_controller)
-
-    def test_register_controller_no_get_info(self):
-        mock_test_config = self.base_mock_test_config.copy()
-        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-        get_info = getattr(mock_controller, 'get_info')
-        delattr(mock_controller, 'get_info')
-        try:
-            mock_test_config.controller_configs = {
-                mock_ctrlr_config_name: ['magic1', 'magic2']
-            }
-            tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-            tr._register_controller(mock_test_config, mock_controller)
-            self.assertEqual(tr.results.controller_info, {})
-        finally:
-            setattr(mock_controller, 'get_info', get_info)
-
-    def test_register_controller_return_value(self):
-        mock_test_config = self.base_mock_test_config.copy()
-        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-        mock_test_config.controller_configs = {
-            mock_ctrlr_config_name: ['magic1', 'magic2']
-        }
-        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-        magic_devices = tr._register_controller(mock_test_config,
-                                                mock_controller)
-        self.assertEqual(magic_devices[0].magic, 'magic1')
-        self.assertEqual(magic_devices[1].magic, 'magic2')
-
-    def test_register_controller_change_return_value(self):
-        mock_test_config = self.base_mock_test_config.copy()
-        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-        mock_test_config.controller_configs = {
-            mock_ctrlr_config_name: ['magic1', 'magic2']
-        }
-        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-        magic_devices = tr._register_controller(mock_test_config,
-                                                mock_controller)
-        magic1 = magic_devices.pop(0)
-        self.assertIs(magic1, tr._controller_registry['mock_controller'][0])
-        self.assertEqual(len(tr._controller_registry['mock_controller']), 2)
-
-    def test_register_controller_less_than_min_number(self):
-        mock_test_config = self.base_mock_test_config.copy()
-        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-        mock_test_config.controller_configs = {
-            mock_ctrlr_config_name: ['magic1', 'magic2']
-        }
-        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
-        expected_msg = 'Expected to get at least 3 controller objects, got 2.'
-        with self.assertRaisesRegex(signals.ControllerError, expected_msg):
-            tr._register_controller(
-                mock_test_config, mock_controller, min_number=3)
-
     def test_run_twice(self):
         """Verifies that:
         1. Repeated run works properly.
@@ -162,13 +74,9 @@ class TestRunnerTest(unittest.TestCase):
         tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
         tr.add_test_class(mock_test_config, integration_test.IntegrationTest)
         tr.run()
-        self.assertFalse(tr._controller_registry)
-        self.assertFalse(tr._controller_destructors)
         self.assertTrue(
             mock_test_config.controller_configs[mock_ctrlr_config_name][0])
         tr.run()
-        self.assertFalse(tr._controller_registry)
-        self.assertFalse(tr._controller_destructors)
         results = tr.results.summary_dict()
         self.assertEqual(results['Requested'], 2)
         self.assertEqual(results['Executed'], 2)
@@ -253,8 +161,6 @@ class TestRunnerTest(unittest.TestCase):
         tr.add_test_class(mock_test_config, integration2_test.Integration2Test)
         tr.add_test_class(mock_test_config, integration_test.IntegrationTest)
         tr.run()
-        self.assertFalse(tr._controller_registry)
-        self.assertFalse(tr._controller_destructors)
         results = tr.results.summary_dict()
         self.assertEqual(results['Requested'], 2)
         self.assertEqual(results['Executed'], 2)
@@ -342,29 +248,6 @@ class TestRunnerTest(unittest.TestCase):
         with self.assertRaisesRegex(test_runner.Error, 'No tests to execute.'):
             tr.run()
 
-    def test_verify_controller_module(self):
-        test_runner.verify_controller_module(mock_controller)
-
-    def test_verify_controller_module_null_attr(self):
-        try:
-            tmp = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-            mock_controller.MOBLY_CONTROLLER_CONFIG_NAME = None
-            msg = 'Controller interface .* in .* cannot be null.'
-            with self.assertRaisesRegex(signals.ControllerError, msg):
-                test_runner.verify_controller_module(mock_controller)
-        finally:
-            mock_controller.MOBLY_CONTROLLER_CONFIG_NAME = tmp
-
-    def test_verify_controller_module_missing_attr(self):
-        try:
-            tmp = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
-            delattr(mock_controller, 'MOBLY_CONTROLLER_CONFIG_NAME')
-            msg = 'Module .* missing required controller module attribute'
-            with self.assertRaisesRegex(signals.ControllerError, msg):
-                test_runner.verify_controller_module(mock_controller)
-        finally:
-            setattr(mock_controller, 'MOBLY_CONTROLLER_CONFIG_NAME', tmp)
-
     @mock.patch(
         'mobly.test_runner._find_test_class',
         return_value=type('SampleTest', (), {}))
diff --git a/tests/mobly/test_suite_test.py b/tests/mobly/test_suite_test.py
index 7e15f25..9b1efc1 100755
--- a/tests/mobly/test_suite_test.py
+++ b/tests/mobly/test_suite_test.py
@@ -55,30 +55,24 @@ class TestSuiteTest(unittest.TestCase):
     def tearDown(self):
         shutil.rmtree(self.tmp_dir)
 
-    def test_controller_object_not_persistent_across_classes_in_the_same_run(
-            self):
-        self.foo_test_controller_obj_id = None
-        self.bar_test_controller_obj_id = None
+    def test_controller_object_not_persistent_across_classes(self):
         test_run_config = self.base_mock_test_config.copy()
         test_run_config.controller_configs = {'MagicDevice': [{'serial': 1}]}
 
         class FooTest(base_test.BaseTestClass):
-            def setup_class(cls):
-                cls.controller = cls.register_controller(mock_controller)[0]
-                self.foo_test_controller_obj_id = id(cls.controller)
+            def setup_class(cls1):
+                self.controller1 = cls1.register_controller(mock_controller)[0]
 
         class BarTest(base_test.BaseTestClass):
-            def setup_class(cls):
-                cls.controller = cls.register_controller(mock_controller)[0]
-                self.bar_test_controller_obj_id = id(cls.controller)
+            def setup_class(cls2):
+                self.controller2 = cls2.register_controller(mock_controller)[0]
 
         tr = test_runner.TestRunner(self.tmp_dir,
                                     test_run_config.test_bed_name)
         tr.add_test_class(test_run_config, FooTest)
         tr.add_test_class(test_run_config, BarTest)
         tr.run()
-        self.assertNotEqual(self.foo_test_controller_obj_id,
-                            self.bar_test_controller_obj_id)
+        self.assertIsNot(self.controller1, self.controller2)
 
 
 if __name__ == "__main__":