Create a proper manager for controllers. (#497)

Moving existing controller management logic to a dedicated manager class.
This improves code quality and reduces maintain cost.

Base class uses calls APIs in the new controller manager instead of holding all the controller states itself directly.
diff --git a/mobly/base_test.py b/mobly/base_test.py
index 045fee8..5dc5e8f 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -21,6 +21,7 @@
 
 from future.utils import raise_with_traceback
 
+from mobly import controller_manager
 from mobly import expects
 from mobly import records
 from mobly import runtime_test_info
@@ -46,30 +47,6 @@
     """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.
 
@@ -97,8 +74,6 @@
         log_path: string, specifies the root directory for all logs written
             by a test run.
         test_bed_name: string, the name of the test bed used by a test run.
-        controller_configs: dict, configs used for instantiating controller
-            objects.
         user_params: dict, custom parameters from user, to be consumed by
             the test logic.
     """
@@ -125,7 +100,6 @@
             self.TAG = self._class_name
         # Set params.
         self.log_path = configs.log_path
-        self.controller_configs = configs.controller_configs
         self.test_bed_name = configs.test_bed_name
         self.user_params = configs.user_params
         self.results = records.TestResult()
@@ -133,10 +107,8 @@
         # Deprecated, use `self.current_test_info.name`.
         self.current_test_name = None
         self._generated_test_table = collections.OrderedDict()
-        # Controller object management.
-        self._controller_registry = collections.OrderedDict(
-        )  # controller_name: objects
-        self._controller_modules = {}  # controller_name: module
+        self._controller_manager = controller_manager.ControllerManager(
+            class_name=self.TAG, controller_configs=configs.controller_configs)
 
     def __enter__(self):
         return self
@@ -269,86 +241,15 @@
                 * `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)
-        logging.debug('Found %d objects for controller %s', len(objects),
-                      module_config_name)
-        self._controller_modules[module_ref_name] = module
-        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, module in self._controller_modules.items():
-            try:
-                logging.debug('Destroying %s.', name)
-                module.destroy(self._controller_registry[name])
-            except:
-                logging.exception('Exception occurred destroying %s.', name)
-        self._controller_registry = collections.OrderedDict()
-        self._controller_modules = {}
+        return self._controller_manager.register_controller(
+            module, required, min_number)
 
     def _record_controller_info(self):
         # Collect controller information and write to test result.
-        for module_ref_name, objects in self._controller_registry.items():
-            module = self._controller_modules[module_ref_name]
-            try:
-                controller_info = module.get_info(copy.copy(objects))
-            except AttributeError:
-                logging.warning('No optional debug info found for controller '
-                                '%s. To provide it, implement `get_info`.',
-                                module_ref_name)
-                continue
-            self.results.add_controller_info(
-                controller_name=module.MOBLY_CONTROLLER_CONFIG_NAME,
-                controller_info=controller_info,
-                test_class=self.TAG)
+        for record in self._controller_manager.get_controller_info_records():
+            self.results.add_controller_info_record(record)
+            self.summary_writer.dump(
+                record.to_dict(), records.TestSummaryEntryType.CONTROLLER_INFO)
 
     def _setup_generated_tests(self):
         """Proxy function to guarantee the base implementation of
@@ -423,6 +324,10 @@
             self.results.add_class_error(record)
             self.summary_writer.dump(record.to_dict(),
                                      records.TestSummaryEntryType.RECORD)
+        finally:
+            # Write controller info and summary to summary file.
+            self._record_controller_info()
+            self._controller_manager.unregister_controllers()
 
     def teardown_class(self):
         """Teardown function that will be called after all the selected tests in
@@ -905,14 +810,7 @@
             setattr(e, 'results', self.results)
             raise e
         finally:
-            # Write controller info and summary to summary file.
-            self._record_controller_info()
-            for controller_info in self.results.controller_info:
-                self.summary_writer.dump(
-                    controller_info.to_dict(),
-                    records.TestSummaryEntryType.CONTROLLER_INFO)
             self._teardown_class()
-            self._unregister_controllers()
             logging.info('Summary for test class %s: %s', self.TAG,
                          self.results.summary_str())
 
diff --git a/mobly/controller_manager.py b/mobly/controller_manager.py
new file mode 100644
index 0000000..db23688
--- /dev/null
+++ b/mobly/controller_manager.py
@@ -0,0 +1,207 @@
+# Copyright 2018 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.
+""" Module for Mobly controller management."""
+import collections
+import copy
+import logging
+import yaml
+
+from mobly import records
+from mobly import signals
+
+
+def verify_controller_module(module):
+    """Verifies a module object follows the required interface for
+    controllers.
+
+    The interface is explained in the docstring of
+    `base_test.BaseTestClass.register_controller`.
+
+    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 ControllerManager(object):
+    """Manages the controller objects for Mobly.
+
+    This manages the life cycles and info retrieval of all controller objects
+    used in a test.
+    """
+
+    def __init__(self, class_name, controller_configs):
+        # Controller object management.
+        self._controller_objects = collections.OrderedDict(
+        )  # controller_name: objects
+        self._controller_modules = {}  # controller_name: module
+        self._class_name = class_name
+        self._controller_configs = controller_configs
+
+    def register_controller(self, module, required=True, min_number=1):
+        """Loads a controller module and returns its loaded devices.
+
+        This is to be used in a mobly test class.
+
+        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_objects:
+            raise signals.ControllerError(
+                'Controller module %s has already been registered. It cannot '
+                'be registered again.' % module_ref_name)
+        # Create controller objects.
+        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 = module.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_objects[module_ref_name] = copy.copy(objects)
+        logging.debug('Found %d objects for controller %s', len(objects),
+                      module_config_name)
+        self._controller_modules[module_ref_name] = module
+        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, module in self._controller_modules.items():
+            logging.debug('Destroying %s.', name)
+            try:
+                module.destroy(self._controller_objects[name])
+            except:
+                logging.exception('Exception occurred destroying %s.', name)
+        self._controller_objects = collections.OrderedDict()
+        self._controller_modules = {}
+
+    def _create_controller_info_record(self, controller_module_name):
+        """Creates controller info record for a particular controller type.
+
+        Info is retrieved from all the controller objects spawned from the
+        specified module, using the controller module's `get_info` function.
+
+        Args:
+            controller_module_name: string, the name of the controller module
+                to retrieve info from.
+
+        Returns:
+            A records.ControllerInfoRecord object.
+        """
+        module = self._controller_modules[controller_module_name]
+        controller_info = None
+        try:
+            controller_info = module.get_info(
+                copy.copy(self._controller_objects[controller_module_name]))
+        except AttributeError:
+            logging.warning('No optional debug info found for controller '
+                            '%s. To provide it, implement `get_info`.',
+                            controller_module_name)
+        try:
+            yaml.dump(controller_info)
+        except TypeError:
+            logging.warning('The info of controller %s in class "%s" is not '
+                            'YAML serializable! Coercing it to string.',
+                            controller_module_name, self._class_name)
+            controller_info = str(controller_info)
+        return records.ControllerInfoRecord(
+            self._class_name, module.MOBLY_CONTROLLER_CONFIG_NAME,
+            controller_info)
+
+    def get_controller_info_records(self):
+        """Get the info records for all the controller objects in the manager.
+
+        New info records for each controller object are created for every call
+        so the latest info is included.
+
+        Returns:
+            List of records.ControllerInfoRecord objects. Each opject conatins
+            the info of a type of controller
+        """
+        info_records = []
+        for controller_module_name in self._controller_objects.keys():
+            record = self._create_controller_info_record(
+                controller_module_name)
+            if record:
+                info_records.append(record)
+        return info_records
diff --git a/mobly/records.py b/mobly/records.py
index 171c0ca..1dec134 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -539,28 +539,16 @@
         else:
             self.error.append(record)
 
-    def add_controller_info(self, controller_name, controller_info,
-                            test_class):
-        """Adds controller info to results.
+    def add_controller_info_record(self, controller_info_record):
+        """Adds a controller info record to results.
 
-        This can be called multiple times for each
+        This can be called multiple times for each test class.
 
         Args:
-            controller_name: string, name of the controller.
-            controller_info: yaml serializable info about the controller.
-            test_class: string, a tag for identifying a class. This should be
-                the test class's own `TAG` attribute.
+            controller_info_record: ControllerInfoRecord object to be added to
+                the result.
         """
-        info = controller_info
-        try:
-            yaml.dump(controller_info)
-        except TypeError:
-            logging.warning('The info of controller %s in class "%s" is not '
-                            'YAML serializable! Coercing it to string.',
-                            controller_name, test_class)
-            info = str(controller_info)
-        self.controller_info.append(
-            ControllerInfoRecord(test_class, controller_name, info))
+        self.controller_info.append(controller_info_record)
 
     def add_class_error(self, test_record):
         """Add a record to indicate a test class has failed before any test
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index f6349ea..ff866bb 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -267,14 +267,25 @@
         on_fail_call_check.assert_called_once_with("haha")
 
     def test_teardown_class_fail_by_exception(self):
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_ctrlr_2_config_name = mock_second_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        my_config = [{'serial': 'xxxx', 'magic': 'Magic'}]
+        mock_test_config.controller_configs[mock_ctrlr_config_name] = my_config
+        mock_test_config.controller_configs[
+            mock_ctrlr_2_config_name] = copy.copy(my_config)
+
         class MockBaseTest(base_test.BaseTestClass):
+            def setup_class(self):
+                self.register_controller(mock_controller)
+
             def test_something(self):
                 pass
 
             def teardown_class(self):
                 raise Exception(MSG_EXPECTED_EXCEPTION)
 
-        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls = MockBaseTest(mock_test_config)
         bt_cls.run()
         test_record = bt_cls.results.passed[0]
         class_record = bt_cls.results.error[0]
@@ -287,6 +298,53 @@
         expected_summary = ('Error 1, Executed 1, Failed 0, Passed 1, '
                             'Requested 1, Skipped 0')
         self.assertEqual(bt_cls.results.summary_str(), expected_summary)
+        # Verify the controller info is recorded correctly.
+        info = bt_cls.results.controller_info[0]
+        self.assertEqual(info.test_class, 'MockBaseTest')
+        self.assertEqual(info.controller_name, 'MagicDevice')
+        self.assertEqual(info.controller_info, [{
+            'MyMagic': {
+                'magic': 'Magic'
+            }
+        }])
+
+    def test_teardown_class_raise_abort_all(self):
+        mock_test_config = self.mock_test_cls_configs.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_ctrlr_2_config_name = mock_second_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        my_config = [{'serial': 'xxxx', 'magic': 'Magic'}]
+        mock_test_config.controller_configs[mock_ctrlr_config_name] = my_config
+        mock_test_config.controller_configs[
+            mock_ctrlr_2_config_name] = copy.copy(my_config)
+
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_class(self):
+                self.register_controller(mock_controller)
+
+            def test_something(self):
+                pass
+
+            def teardown_class(self):
+                raise asserts.abort_all(MSG_EXPECTED_EXCEPTION)
+
+        bt_cls = MockBaseTest(mock_test_config)
+        with self.assertRaisesRegex(signals.TestAbortAll,
+                                    MSG_EXPECTED_EXCEPTION):
+            bt_cls.run()
+        test_record = bt_cls.results.passed[0]
+        self.assertTrue(bt_cls.results.is_all_pass)
+        expected_summary = ('Error 0, Executed 1, Failed 0, Passed 1, '
+                            'Requested 1, Skipped 0')
+        self.assertEqual(bt_cls.results.summary_str(), expected_summary)
+        # Verify the controller info is recorded correctly.
+        info = bt_cls.results.controller_info[0]
+        self.assertEqual(info.test_class, 'MockBaseTest')
+        self.assertEqual(info.controller_name, 'MagicDevice')
+        self.assertEqual(info.controller_info, [{
+            'MyMagic': {
+                'magic': 'Magic'
+            }
+        }])
 
     def test_setup_test_fail_by_exception(self):
         mock_on_fail = mock.Mock()
@@ -1836,113 +1894,6 @@
             }
         }])
 
-    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_modules[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/controller_manager_test.py b/tests/mobly/controller_manager_test.py
new file mode 100755
index 0000000..7c9f1b6
--- /dev/null
+++ b/tests/mobly/controller_manager_test.py
@@ -0,0 +1,206 @@
+# Copyright 2018 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.
+"""Unit tests for controller manager."""
+import mock
+
+from future.tests.base import unittest
+
+from mobly import controller_manager
+from mobly import signals
+
+from tests.lib import mock_controller
+
+
+class ControllerManagerTest(unittest.TestCase):
+    """Unit tests for Mobly's ControllerManager."""
+
+    def test_verify_controller_module(self):
+        controller_manager.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):
+                controller_manager.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):
+                controller_manager.verify_controller_module(mock_controller)
+        finally:
+            setattr(mock_controller, 'MOBLY_CONTROLLER_CONFIG_NAME', tmp)
+
+    def test_register_controller_no_config(self):
+        c_manager = controller_manager.ControllerManager('SomeClass', {})
+        with self.assertRaisesRegex(signals.ControllerError,
+                                    'No corresponding config found for'):
+            c_manager.register_controller(mock_controller)
+
+    def test_register_controller_no_config_for_not_required(self):
+        c_manager = controller_manager.ControllerManager('SomeClass', {})
+        self.assertIsNone(
+            c_manager.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_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        c_manager.register_controller(mock_controller)
+        registered_name = 'mock_controller'
+        self.assertTrue(registered_name in c_manager._controller_objects)
+        mock_ctrlrs = c_manager._controller_objects[registered_name]
+        self.assertEqual(mock_ctrlrs[0].magic, 'magic1')
+        self.assertEqual(mock_ctrlrs[1].magic, 'magic2')
+        self.assertTrue(c_manager._controller_modules[registered_name])
+        expected_msg = 'Controller module .* has already been registered.'
+        with self.assertRaisesRegex(signals.ControllerError, expected_msg):
+            c_manager.register_controller(mock_controller)
+
+    def test_register_controller_return_value(self):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        magic_devices = c_manager.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_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        magic_devices = c_manager.register_controller(mock_controller)
+        magic1 = magic_devices.pop(0)
+        self.assertIs(magic1,
+                      c_manager._controller_objects['mock_controller'][0])
+        self.assertEqual(
+            len(c_manager._controller_objects['mock_controller']), 2)
+
+    def test_register_controller_less_than_min_number(self):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        expected_msg = 'Expected to get at least 3 controller objects, got 2.'
+        with self.assertRaisesRegex(signals.ControllerError, expected_msg):
+            c_manager.register_controller(mock_controller, min_number=3)
+
+    @mock.patch('yaml.dump', side_effect=TypeError('ha'))
+    def test_get_controller_info_record_not_serializable(self, _):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        c_manager.register_controller(mock_controller)
+        record = c_manager.get_controller_info_records()[0]
+        actual_controller_info = record.controller_info
+        self.assertEqual(actual_controller_info,
+                         "[{'MyMagic': 'magic1'}, {'MyMagic': 'magic2'}]")
+
+    def test_controller_record_exists_without_get_info(self):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        get_info = getattr(mock_controller, 'get_info')
+        delattr(mock_controller, 'get_info')
+        try:
+            c_manager.register_controller(mock_controller)
+            record = c_manager.get_controller_info_records()[0]
+            self.assertIsNone(record.controller_info)
+            self.assertEqual(record.test_class, 'SomeClass')
+            self.assertEqual(record.controller_name, 'MagicDevice')
+        finally:
+            setattr(mock_controller, 'get_info', get_info)
+
+    @mock.patch('tests.lib.mock_controller.get_info')
+    def test_get_controller_info_records_empty(self, mock_get_info_func):
+        mock_get_info_func.return_value = None
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        c_manager.register_controller(mock_controller)
+        record = c_manager.get_controller_info_records()[0]
+        self.assertIsNone(record.controller_info)
+        self.assertEqual(record.test_class, 'SomeClass')
+        self.assertEqual(record.controller_name, 'MagicDevice')
+
+    def test_get_controller_info_records(self):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        c_manager.register_controller(mock_controller)
+        record = c_manager.get_controller_info_records()[0]
+        record_dict = record.to_dict()
+        record_dict.pop('Timestamp')
+        self.assertEqual(
+            record_dict, {
+                'Controller Info': [{
+                    'MyMagic': 'magic1'
+                }, {
+                    'MyMagic': 'magic2'
+                }],
+                'Controller Name': 'MagicDevice',
+                'Test Class': 'SomeClass'
+            })
+
+    def test_get_controller_info_without_registration(self):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        self.assertFalse(c_manager.get_controller_info_records())
+
+    @mock.patch('tests.lib.mock_controller.destroy')
+    def test_unregister_controller(self, mock_destroy_func):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        objects = c_manager.register_controller(mock_controller)
+        c_manager.unregister_controllers()
+        mock_destroy_func.assert_called_once_with(objects)
+        self.assertFalse(c_manager._controller_objects)
+        self.assertFalse(c_manager._controller_modules)
+
+    @mock.patch('tests.lib.mock_controller.destroy')
+    def test_unregister_controller_without_registration(
+            self, mock_destroy_func):
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        controller_configs = {mock_ctrlr_config_name: ['magic1', 'magic2']}
+        c_manager = controller_manager.ControllerManager(
+            'SomeClass', controller_configs)
+        c_manager.unregister_controllers()
+        mock_destroy_func.assert_not_called()
+        self.assertFalse(c_manager._controller_objects)
+        self.assertFalse(c_manager._controller_modules)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index a967b8b..9648637 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -239,15 +239,18 @@
         record1.test_pass(s)
         tr1 = records.TestResult()
         tr1.add_record(record1)
-        tr1.add_controller_info('SomeClass', 'MockDevice',
-                                ['magicA', 'magicB'])
+        controller_info = records.ControllerInfoRecord(
+            'SomeClass', 'MockDevice', ['magicA', 'magicB'])
+        tr1.add_controller_info_record(controller_info)
         record2 = records.TestResultRecord(self.tn)
         record2.test_begin()
         s = signals.TestPass(self.details, self.json_extra)
         record2.test_pass(s)
         tr2 = records.TestResult()
         tr2.add_record(record2)
-        tr2.add_controller_info('SomeClass', 'MockDevice', ['magicC'])
+        controller_info = records.ControllerInfoRecord(
+            'SomeClass', 'MockDevice', ['magicC'])
+        tr2.add_controller_info_record(controller_info)
         tr2 += tr1
         self.assertTrue(tr2.passed, [tr1, tr2])
         self.assertTrue(tr2.controller_info, {'MockDevice': ['magicC']})
@@ -413,25 +416,17 @@
         self.assertIsNot(er, new_er)
         self.assertDictEqual(er.to_dict(), new_er.to_dict())
 
-    def test_add_controller_info(self):
+    def test_add_controller_info_record(self):
         tr = records.TestResult()
         self.assertFalse(tr.controller_info)
-        tr.add_controller_info('MockDevice', ['magicA', 'magicB'], 'MyTest')
+        controller_info = records.ControllerInfoRecord(
+            'SomeClass', 'MockDevice', ['magicA', 'magicB'])
+        tr.add_controller_info_record(controller_info)
         self.assertTrue(tr.controller_info[0])
         self.assertEqual(tr.controller_info[0].controller_name, 'MockDevice')
         self.assertEqual(tr.controller_info[0].controller_info,
                          ['magicA', 'magicB'])
 
-    @mock.patch('yaml.dump', side_effect=TypeError('ha'))
-    def test_add_controller_info_not_serializable(self, mock_yaml_dump):
-        tr = records.TestResult()
-        self.assertFalse(tr.controller_info)
-        tr.add_controller_info('MockDevice', ['magicA', 'magicB'], 'MyTest')
-        self.assertTrue(tr.controller_info[0])
-        self.assertEqual(tr.controller_info[0].controller_name, 'MockDevice')
-        self.assertEqual(tr.controller_info[0].controller_info,
-                         "['magicA', 'magicB']")
-
 
 if __name__ == '__main__':
     unittest.main()