Properly record controller info. (#495)

* Support multiple groups of controller info in test result record.
* Record controller info at the end instead of the beginning of a test class run.
* Introduce a proper data class for controller info.
* Add the new module needed for the new unit tests.

7 files changed, 254 insertions, 78 deletions

diff --git a/mobly/base_test.py b/mobly/base_test.py
index e4139d9..045fee8 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -134,8 +134,9 @@ class BaseTestClass(object):
         self.current_test_name = None
         self._generated_test_table = collections.OrderedDict()
         # Controller object management.
-        self._controller_registry = {}
-        self._controller_destructors = {}
+        self._controller_registry = collections.OrderedDict(
+        )  # controller_name: objects
+        self._controller_modules = {}  # controller_name: module
 
     def __enter__(self):
         return self
@@ -312,22 +313,9 @@ class BaseTestClass(object):
         # 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
+        self._controller_modules[module_ref_name] = module
         return objects
 
     def _unregister_controllers(self):
@@ -337,14 +325,30 @@ class BaseTestClass(object):
         """
         # TODO(xpconanfan): actually record these errors instead of just
         # logging them.
-        for name, destroy in self._controller_destructors.items():
+        for name, module in self._controller_modules.items():
             try:
                 logging.debug('Destroying %s.', name)
-                destroy(self._controller_registry[name])
+                module.destroy(self._controller_registry[name])
             except:
                 logging.exception('Exception occurred destroying %s.', name)
-        self._controller_registry = {}
-        self._controller_destructors = {}
+        self._controller_registry = collections.OrderedDict()
+        self._controller_modules = {}
+
+    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)
 
     def _setup_generated_tests(self):
         """Proxy function to guarantee the base implementation of
@@ -901,6 +905,12 @@ class BaseTestClass(object):
             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,
diff --git a/mobly/records.py b/mobly/records.py
index 051ff9c..171c0ca 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -21,6 +21,7 @@ import io
 import logging
 import sys
 import threading
+import time
 import traceback
 import yaml
 
@@ -52,7 +53,7 @@ class TestSummaryEntryType(enum.Enum):
     RECORD = 'Record'
     # A summary of the test run stats, e.g. how many test failed.
     SUMMARY = 'Summary'
-    # Information on the controllers used in the test.
+    # Information on the controllers used in a test class.
     CONTROLLER_INFO = 'ControllerInfo'
     # Additional data added by users during test.
     # This can be added at any point in the test, so do not assume the location
@@ -155,8 +156,34 @@ class TestResultEnums(object):
     TEST_RESULT_ERROR = 'ERROR'
 
 
+class ControllerInfoRecord(object):
+    """A record representing the controller info in test results."""
+
+    KEY_TEST_CLASS = TestResultEnums.RECORD_CLASS
+    KEY_CONTROLLER_NAME = 'Controller Name'
+    KEY_CONTROLLER_INFO = 'Controller Info'
+    KEY_TIMESTAMP = 'Timestamp'
+
+    def __init__(self, test_class, controller_name, info):
+        self.test_class = test_class
+        self.controller_name = controller_name
+        self.controller_info = info
+        self.timestamp = time.time()
+
+    def to_dict(self):
+        result = {}
+        result[self.KEY_TEST_CLASS] = self.test_class
+        result[self.KEY_CONTROLLER_NAME] = self.controller_name
+        result[self.KEY_CONTROLLER_INFO] = self.controller_info
+        result[self.KEY_TIMESTAMP] = self.timestamp
+        return result
+
+    def __repr__(self):
+        return str(self.to_dict())
+
+
 class ExceptionRecord(object):
-    """A wrapper class for representing exception objects in TestResultRecord.
+    """A record representing exception objects in TestResultRecord.
 
     Attributes:
         exception: Exception object, the original Exception.
@@ -235,7 +262,7 @@ class ExceptionRecord(object):
         result = ExceptionRecord(exception, self.position)
         result.stacktrace = self.stacktrace
         result.details = self.details
-        result.extras = self.extras
+        result.extras = copy.deepcopy(self.extras)
         result.position = self.position
         return result
 
@@ -447,13 +474,14 @@ class TestResult(object):
     This class is essentially a container of TestResultRecord objects.
 
     Attributes:
-        self.requested: A list of strings, each is the name of a test requested
+        requested: A list of strings, each is the name of a test requested
             by user.
-        self.failed: A list of records for tests failed.
-        self.executed: A list of records for tests that were actually executed.
-        self.passed: A list of records for tests passed.
-        self.skipped: A list of records for tests skipped.
-        self.error: A list of records for tests with error result token.
+        failed: A list of records for tests failed.
+        executed: A list of records for tests that were actually executed.
+        passed: A list of records for tests passed.
+        skipped: A list of records for tests skipped.
+        error: A list of records for tests with error result token.
+        controller_info: list of ControllerInfoRecord.
     """
 
     def __init__(self):
@@ -463,7 +491,7 @@ class TestResult(object):
         self.passed = []
         self.skipped = []
         self.error = []
-        self.controller_info = {}
+        self.controller_info = []
 
     def __add__(self, r):
         """Overrides '+' operator for TestResult class.
@@ -486,12 +514,6 @@ class TestResult(object):
             l_value = getattr(self, name)
             if isinstance(r_value, list):
                 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, use the
-                # r-value which is more up to date.
-                # TODO(xpconanfan): have a better way to validate this situation.
-                setattr(sum_result, name, r_value)
         return sum_result
 
     def add_record(self, record):
@@ -517,15 +539,28 @@ class TestResult(object):
         else:
             self.error.append(record)
 
-    def add_controller_info(self, name, info):
+    def add_controller_info(self, controller_name, controller_info,
+                            test_class):
+        """Adds controller info to results.
+
+        This can be called multiple times for each
+
+        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.
+        """
+        info = controller_info
         try:
-            yaml.dump(info)
+            yaml.dump(controller_info)
         except TypeError:
-            logging.warning('Controller info for %s is not YAML serializable!'
-                            ' Coercing it to string.' % name)
-            self.controller_info[name] = str(info)
-            return
-        self.controller_info[name] = info
+            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))
 
     def add_class_error(self, test_record):
         """Add a record to indicate a test class has failed before any test
diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index 9c57715..909385b 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -351,9 +351,6 @@ class TestRunner(object):
                         'Abort all subsequent test classes. Reason: %s', e)
                     raise
         finally:
-            # Write controller info and summary to summary file.
-            summary_writer.dump(self.results.controller_info,
-                                records.TestSummaryEntryType.CONTROLLER_INFO)
             summary_writer.dump(self.results.summary_dict(),
                                 records.TestSummaryEntryType.SUMMARY)
             # Stop and show summary.
diff --git a/tests/lib/mock_second_controller.py b/tests/lib/mock_second_controller.py
new file mode 100644
index 0000000..a4a847a
--- /dev/null
+++ b/tests/lib/mock_second_controller.py
@@ -0,0 +1,59 @@
+# Copyright 2016 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.
+
+# This is a second mock third-party controller module used for testing Mobly's
+# handling of multiple controller modules.
+
+import logging
+
+MOBLY_CONTROLLER_CONFIG_NAME = "AnotherMagicDevice"
+
+
+def create(configs):
+    objs = []
+    for c in configs:
+        if isinstance(c, dict):
+            c.pop("serial")
+        objs.append(AnotherMagicDevice(c))
+    return objs
+
+
+def destroy(objs):
+    print("Destroying other magic")
+
+
+def get_info(objs):
+    infos = []
+    for obj in objs:
+        infos.append(obj.who_am_i())
+    return infos
+
+
+class AnotherMagicDevice(object):
+    """This controller supports adding controller's info during test.
+
+    It is used for testing that this info is correctly recorded by Mobly.
+    """
+    def __init__(self, config):
+        self.magic = config
+
+    def get_magic(self):
+        logging.info("My other magic is %s.", self.magic)
+        return self.magic
+
+    def set_magic(self, extra_magic):
+        self.magic['extra_magic'] = extra_magic
+
+    def who_am_i(self):
+        return {"MyOtherMagic": self.magic}
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index 898ab5d..ead1bf7 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import io
 import os
 import mock
@@ -30,6 +31,7 @@ from mobly import signals
 
 from tests.lib import utils
 from tests.lib import mock_controller
+from tests.lib import mock_second_controller
 
 MSG_EXPECTED_EXCEPTION = "This is an expected exception."
 MSG_EXPECTED_TEST_FAILURE = "This is an expected test failure."
@@ -1784,6 +1786,56 @@ class BaseTestTest(unittest.TestCase):
                 self.assertIsNotNone(c['timestamp'])
         self.assertTrue(hit)
 
+    def test_record_controller_info(self):
+        """Verifies that controller info is correctly recorded.
+
+        1. Info added in test is recorded.
+        2. Info of multiple controller types are recorded.
+        """
+        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 ControllerInfoTest(base_test.BaseTestClass):
+            """Registers two different controller types and modifies controller
+            info at runtime.
+            """
+
+            def setup_class(self):
+                self.register_controller(mock_controller)
+                second_controller = self.register_controller(
+                    mock_second_controller)[0]
+                # This should appear in recorded controller info.
+                second_controller.set_magic('haha')
+
+            def test_func(self):
+                pass
+
+        bt_cls = ControllerInfoTest(mock_test_config)
+        bt_cls.run()
+        info1 = bt_cls.results.controller_info[0]
+        info2 = bt_cls.results.controller_info[1]
+        self.assertNotEqual(info1.timestamp, info2.timestamp)
+        self.assertEqual(info1.test_class, 'ControllerInfoTest')
+        self.assertEqual(info1.controller_name, 'MagicDevice')
+        self.assertEqual(info1.controller_info, [{
+            'MyMagic': {
+                'magic': 'Magic'
+            }
+        }])
+        self.assertEqual(info2.test_class, 'ControllerInfoTest')
+        self.assertEqual(info2.controller_name, 'AnotherMagicDevice')
+        self.assertEqual(info2.controller_info, [{
+            'MyOtherMagic': {
+                'magic': 'Magic',
+                'extra_magic': 'haha'
+            }
+        }])
+
     def test_register_controller_no_config(self):
         bt_cls = MockEmptyBaseTest(self.mock_test_cls_configs)
         with self.assertRaisesRegex(signals.ControllerError,
@@ -1812,7 +1864,7 @@ class BaseTestTest(unittest.TestCase):
         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])
+        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)
@@ -1828,7 +1880,7 @@ class BaseTestTest(unittest.TestCase):
             }
             bt_cls = MockEmptyBaseTest(mock_test_config)
             bt_cls.register_controller(mock_controller)
-            self.assertEqual(bt_cls.results.controller_info, {})
+            self.assertEqual(bt_cls.results.controller_info, [])
         finally:
             setattr(mock_controller, 'get_info', get_info)
 
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index f818e62..a967b8b 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -1,13 +1,13 @@
 # Copyright 2016 Google Inc.
 #
-# Licensed under the Apache License, Version 2.0 (the "License");
+# 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,
+# 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.
@@ -46,10 +46,10 @@ class RecordsTest(unittest.TestCase):
     """
 
     def setUp(self):
-        self.tn = "test_name"
-        self.details = "Some details about the test execution."
+        self.tn = 'test_name'
+        self.details = 'Some details about the test execution.'
         self.float_extra = 12345.56789
-        self.json_extra = {"ha": "whatever"}
+        self.json_extra = {'ha': 'whatever'}
         self.tmp_path = tempfile.mkdtemp()
 
     def tearDown(self):
@@ -62,8 +62,8 @@ class RecordsTest(unittest.TestCase):
         self.assertEqual(record.result, result)
         self.assertEqual(record.details, details)
         self.assertEqual(record.extras, extras)
-        self.assertTrue(record.begin_time, "begin time should not be empty.")
-        self.assertTrue(record.end_time, "end time should not be empty.")
+        self.assertTrue(record.begin_time, 'begin time should not be empty.')
+        self.assertTrue(record.end_time, 'end time should not be empty.')
         # UID is not used at the moment, should always be None.
         self.assertIsNone(record.uid)
         # Verify to_dict.
@@ -88,7 +88,7 @@ class RecordsTest(unittest.TestCase):
         self.assertDictEqual(actual_d, d)
         # Verify that these code paths do not cause crashes and yield non-empty
         # results.
-        self.assertTrue(str(record), "str of the record should not be empty.")
+        self.assertTrue(str(record), 'str of the record should not be empty.')
         self.assertTrue(repr(record), "the record's repr shouldn't be empty.")
 
     """ Begin of Tests """
@@ -239,17 +239,18 @@ class RecordsTest(unittest.TestCase):
         record1.test_pass(s)
         tr1 = records.TestResult()
         tr1.add_record(record1)
-        tr1.add_controller_info("MockDevice", ["magicA", "magicB"])
+        tr1.add_controller_info('SomeClass', 'MockDevice',
+                                ['magicA', 'magicB'])
         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("MockDevice", ["magicC"])
+        tr2.add_controller_info('SomeClass', 'MockDevice', ['magicC'])
         tr2 += tr1
         self.assertTrue(tr2.passed, [tr1, tr2])
-        self.assertTrue(tr2.controller_info, {"MockDevice": ["magicC"]})
+        self.assertTrue(tr2.controller_info, {'MockDevice': ['magicC']})
 
     def test_result_add_operator_type_mismatch(self):
         record1 = records.TestResultRecord(self.tn)
@@ -258,9 +259,9 @@ class RecordsTest(unittest.TestCase):
         record1.test_pass(s)
         tr1 = records.TestResult()
         tr1.add_record(record1)
-        expected_msg = "Operand .* of type .* is not a TestResult."
+        expected_msg = 'Operand .* of type .* is not a TestResult.'
         with self.assertRaisesRegex(TypeError, expected_msg):
-            tr1 += "haha"
+            tr1 += 'haha'
 
     def test_result_add_class_error_with_test_signal(self):
         record1 = records.TestResultRecord(self.tn)
@@ -270,7 +271,7 @@ class RecordsTest(unittest.TestCase):
         tr = records.TestResult()
         tr.add_record(record1)
         s = signals.TestFailure(self.details, self.float_extra)
-        record2 = records.TestResultRecord("SomeTest", s)
+        record2 = records.TestResultRecord('SomeTest', s)
         tr.add_class_error(record2)
         self.assertEqual(len(tr.passed), 1)
         self.assertEqual(len(tr.error), 1)
@@ -289,10 +290,10 @@ class RecordsTest(unittest.TestCase):
 
         class SpecialError(Exception):
             def __init__(self, arg1, arg2):
-                self.msg = "%s %s" % (arg1, arg2)
+                self.msg = '%s %s' % (arg1, arg2)
 
-        se = SpecialError("haha", 42)
-        record2 = records.TestResultRecord("SomeTest", se)
+        se = SpecialError('haha', 42)
+        record2 = records.TestResultRecord('SomeTest', se)
         tr.add_class_error(record2)
         self.assertEqual(len(tr.passed), 1)
         self.assertEqual(len(tr.error), 1)
@@ -334,7 +335,7 @@ class RecordsTest(unittest.TestCase):
         """
         record1 = records.TestResultRecord(self.tn)
         record1.test_begin()
-        record1.test_fail(Exception("haha"))
+        record1.test_fail(Exception('haha'))
         tr = records.TestResult()
         tr.add_class_error(record1)
         self.assertFalse(tr.is_all_pass)
@@ -342,7 +343,7 @@ class RecordsTest(unittest.TestCase):
     def test_is_test_executed(self):
         record1 = records.TestResultRecord(self.tn)
         record1.test_begin()
-        record1.test_fail(Exception("haha"))
+        record1.test_fail(Exception('haha'))
         tr = records.TestResult()
         tr.add_record(record1)
         self.assertTrue(tr.is_test_executed(record1.test_name))
@@ -415,20 +416,22 @@ class RecordsTest(unittest.TestCase):
     def test_add_controller_info(self):
         tr = records.TestResult()
         self.assertFalse(tr.controller_info)
-        tr.add_controller_info('MockDevice', ['magicA', 'magicB'])
-        self.assertTrue(tr.controller_info)
-        self.assertEqual(tr.controller_info['MockDevice'],
+        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'])
 
     @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'])
-        self.assertTrue(tr.controller_info)
-        self.assertEqual(tr.controller_info['MockDevice'],
+        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__":
+if __name__ == '__main__':
     unittest.main()
diff --git a/tests/mobly/test_runner_test.py b/tests/mobly/test_runner_test.py
index d86f60c..9459f44 100755
--- a/tests/mobly/test_runner_test.py
+++ b/tests/mobly/test_runner_test.py
@@ -55,6 +55,13 @@ class TestRunnerTest(unittest.TestCase):
     def tearDown(self):
         shutil.rmtree(self.tmp_dir)
 
+    def _assertControllerInfoEqual(self, info, expected_info_dict):
+        self.assertEqual(expected_info_dict['Controller Name'],
+                         info.controller_name)
+        self.assertEqual(expected_info_dict['Test Class'], info.test_class)
+        self.assertEqual(expected_info_dict['Controller Info'],
+                         info.controller_info)
+
     def test_run_twice(self):
         """Verifies that:
         1. Repeated run works properly.
@@ -81,8 +88,8 @@ class TestRunnerTest(unittest.TestCase):
         self.assertEqual(results['Requested'], 2)
         self.assertEqual(results['Executed'], 2)
         self.assertEqual(results['Passed'], 2)
-        expected_info = {
-            'MagicDevice': [{
+        expected_info_dict = {
+            'Controller Info': [{
                 'MyMagic': {
                     'magic': 'Magic1'
                 }
@@ -90,9 +97,18 @@ class TestRunnerTest(unittest.TestCase):
                 'MyMagic': {
                     'magic': 'Magic2'
                 }
-            }]
+            }],
+            'Controller Name':
+            'MagicDevice',
+            'Test Class':
+            'IntegrationTest',
         }
-        self.assertEqual(tr.results.controller_info, expected_info)
+        self._assertControllerInfoEqual(tr.results.controller_info[0],
+                                        expected_info_dict)
+        self._assertControllerInfoEqual(tr.results.controller_info[1],
+                                        expected_info_dict)
+        self.assertNotEqual(tr.results.controller_info[0].timestamp,
+                            tr.results.controller_info[1].timestamp)
 
     def test_summary_file_entries(self):
         """Verifies the output summary's file format.
@@ -123,6 +139,10 @@ class TestRunnerTest(unittest.TestCase):
                          records.TestSummaryEntryType.TEST_NAME_LIST.value)
         self.assertEqual(summary_entries[1]['Type'],
                          records.TestSummaryEntryType.RECORD.value)
+        self.assertEqual(summary_entries[2]['Type'],
+                         records.TestSummaryEntryType.CONTROLLER_INFO.value)
+        self.assertEqual(summary_entries[3]['Type'],
+                         records.TestSummaryEntryType.SUMMARY.value)
 
     @mock.patch(
         'mobly.controllers.android_device_lib.adb.AdbProxy',