From 54f9beeab6d71aa65b9ca6f72c5ec0dfd0b4ef02 Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Mon, 2 Dec 2024 19:49:53 +0800
Subject: [PATCH 01/13] Record suite class name in summary yaml.

---
 mobly/records.py                 | 23 +++++++++
 mobly/suite_runner.py            | 12 ++++-
 tests/mobly/records_test.py      | 11 +++++
 tests/mobly/suite_runner_test.py | 80 ++++++++++++++++++++++++--------
 4 files changed, 106 insertions(+), 20 deletions(-)

diff --git a/mobly/records.py b/mobly/records.py
index 10d3063c..c1487cd5 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -94,6 +94,8 @@ class TestSummaryEntryType(enum.Enum):
   SUMMARY = 'Summary'
   # Information on the controllers used in a test class.
   CONTROLLER_INFO = 'ControllerInfo'
+  # Suite level information.
+  SUITE_INFO = 'SuiteInfo'
   # Additional data added by users during test.
   # This can be added at any point in the test, so do not assume the location
   # of these entries in the summary file.
@@ -724,3 +726,24 @@ def summary_dict(self):
     d['Skipped'] = len(self.skipped)
     d['Error'] = len(self.error)
     return d
+
+
+class SuiteInfoRecord:
+  """A record representing the suite info in test summary."""
+
+  KEY_SUITE_CLASS_NAME = 'Suite Class Name'
+  KEY_TIMESTAMP = 'Timestamp'
+
+  def __init__(self, suite_class_name):
+    self.suite_class_name = suite_class_name
+    self.timestamp = time.time()
+
+  def to_dict(self):
+    result = {}
+    result[self.KEY_SUITE_CLASS_NAME] = self.suite_class_name
+    result[self.KEY_TIMESTAMP] = self.timestamp
+    return result
+
+  def __repr__(self):
+    return str(self.to_dict())
+
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index b0b064ca..795e8421 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -68,11 +68,13 @@ def setup_suite(self, config):
 import collections
 import inspect
 import logging
+import os
 import sys
 
 from mobly import base_test
 from mobly import base_suite
 from mobly import config_parser
+from mobly import records
 from mobly import signals
 from mobly import test_runner
 
@@ -188,6 +190,13 @@ def _print_test_names(test_classes):
       print(f'{cls.TAG}.{name}')
 
 
+def _record_suite_info(log_path, suite_class_name):
+  summary_path = os.path.join(log_path, records.OUTPUT_FILE_SUMMARY)
+  summary_writer = records.TestSummaryWriter(summary_path)
+  content = records.SuiteInfoRecord(suite_class_name=suite_class_name)
+  summary_writer.dump(content.to_dict(), records.TestSummaryEntryType.SUITE_INFO)
+
+
 def run_suite_class(argv=None):
   """Executes tests in the test suite.
 
@@ -212,7 +221,8 @@ def run_suite_class(argv=None):
   suite = suite_class(runner, config)
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
-  with runner.mobly_logger(console_level=console_level):
+  with runner.mobly_logger(console_level=console_level) as log_path:
+    _record_suite_info(log_path, suite_class_name=suite_class.__name__)
     try:
       suite.setup_suite(config.copy())
       try:
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index 5c1d027f..f9648bec 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -542,6 +542,17 @@ def test_uid_helper():
 
     self.assertEqual(test_uid_helper.uid, 'some-uuid')
 
+  def test_create_suite_info_record_and_convert_to_dict(self):
+    suite_class_name = 'FakeTestSuite'
+    record = records.SuiteInfoRecord(suite_class_name=suite_class_name)
+
+    result = record.to_dict()
+
+    self.assertIn(
+        (records.SuiteInfoRecord.KEY_SUITE_CLASS_NAME, suite_class_name),
+        result.items(),
+    )
+
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 08072d6b..99f8af7d 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -13,18 +13,22 @@
 # limitations under the License.
 
 import io
+import logging
 import os
 import shutil
 import sys
 import tempfile
+import time
 import unittest
 from unittest import mock
 
 from mobly import base_suite
 from mobly import base_test
+from mobly import records
 from mobly import suite_runner
 from tests.lib import integration2_test
 from tests.lib import integration_test
+import yaml
 
 
 class FakeTest1(base_test.BaseTestClass):
@@ -136,8 +140,9 @@ def test_run_suite_with_failures(self, mock_exit):
     mock_exit.assert_called_once_with(1)
 
   @mock.patch('sys.exit')
-  @mock.patch.object(suite_runner, '_find_suite_class', autospec=True)
-  def test_run_suite_class(self, mock_find_suite_class, mock_exit):
+  def test_run_suite_class(self, mock_exit):
+    tmp_file_path = self._gen_tmp_config_file()
+    mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
     mock_called = mock.MagicMock()
 
     class FakeTestSuite(base_suite.BaseSuite):
@@ -151,30 +156,53 @@ def teardown_suite(self):
         mock_called.teardown_suite()
         super().teardown_suite()
 
-    mock_find_suite_class.return_value = FakeTestSuite
-
-    tmp_file_path = os.path.join(self.tmp_dir, 'config.yml')
-    with io.open(tmp_file_path, 'w', encoding='utf-8') as f:
-      f.write(
-          """
-        TestBeds:
-          # A test bed where adb will find Android devices.
-          - Name: SampleTestBed
-            Controllers:
-              MagicDevice: '*'
-      """
-      )
-
-    mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
+    sys.modules['__main__'].__dict__[FakeTestSuite.__name__] = FakeTestSuite
 
     with mock.patch.object(sys, 'argv', new=mock_cli_args):
-      suite_runner.run_suite_class()
+      try:
+        suite_runner.run_suite_class()
+      finally:
+        del sys.modules['__main__'].__dict__[FakeTestSuite.__name__]
 
-    mock_find_suite_class.assert_called_once()
     mock_called.setup_suite.assert_called_once_with()
     mock_called.teardown_suite.assert_called_once_with()
     mock_exit.assert_not_called()
 
+  @mock.patch('sys.exit')
+  @mock.patch.object(time, 'time', return_value=1733143236.278318)
+  def test_run_suite_class_records_suite_class_name(self, mock_time, _):
+    tmp_file_path = self._gen_tmp_config_file()
+    mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
+    expected_summary_entry = records.SuiteInfoRecord(
+        suite_class_name='FakeTestSuite'
+    ).to_dict()
+    expected_summary_entry['Type'] = records.TestSummaryEntryType.SUITE_INFO.value
+
+    class FakeTestSuite(base_suite.BaseSuite):
+
+      def setup_suite(self, config):
+        super().setup_suite(config)
+        self.add_test_class(FakeTest1)
+
+    sys.modules['__main__'].__dict__[FakeTestSuite.__name__] = FakeTestSuite
+
+    with mock.patch.object(sys, 'argv', new=mock_cli_args):
+      try:
+        suite_runner.run_suite_class()
+      finally:
+        del sys.modules['__main__'].__dict__[FakeTestSuite.__name__]
+
+    summary_path = os.path.join(
+        logging.root_output_path, records.OUTPUT_FILE_SUMMARY
+    )
+    with io.open(summary_path, 'r', encoding='utf-8') as f:
+      summary_entries = list(yaml.safe_load_all(f))
+
+    self.assertIn(
+        expected_summary_entry,
+        summary_entries,
+    )
+
   def test_print_test_names(self):
     mock_test_class = mock.MagicMock()
     mock_cls_instance = mock.MagicMock()
@@ -191,6 +219,20 @@ def test_print_test_names_with_exception(self):
     mock_cls_instance._pre_run.side_effect = Exception('Something went wrong.')
     mock_cls_instance._clean_up.assert_called_once()
 
+  def _gen_tmp_config_file(self):
+    tmp_file_path = os.path.join(self.tmp_dir, 'config.yml')
+    with io.open(tmp_file_path, 'w', encoding='utf-8') as f:
+      f.write(
+          """
+        TestBeds:
+          # A test bed where adb will find Android devices.
+          - Name: SampleTestBed
+            Controllers:
+              MagicDevice: '*'
+      """
+      )
+    return tmp_file_path
+
 
 if __name__ == '__main__':
   unittest.main()

From fd0c097a297bf132bae107a2dc8f6a925cfbb9ae Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Mon, 2 Dec 2024 21:08:04 +0800
Subject: [PATCH 02/13] Minor fixes and run pyink.

---
 mobly/records.py                 | 3 +--
 mobly/suite_runner.py            | 5 ++++-
 tests/mobly/records_test.py      | 2 +-
 tests/mobly/suite_runner_test.py | 4 +++-
 4 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/mobly/records.py b/mobly/records.py
index c1487cd5..970e2ea6 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -94,7 +94,7 @@ class TestSummaryEntryType(enum.Enum):
   SUMMARY = 'Summary'
   # Information on the controllers used in a test class.
   CONTROLLER_INFO = 'ControllerInfo'
-  # Suite level information.
+  # Test suite level information.
   SUITE_INFO = 'SuiteInfo'
   # Additional data added by users during test.
   # This can be added at any point in the test, so do not assume the location
@@ -746,4 +746,3 @@ def to_dict(self):
 
   def __repr__(self):
     return str(self.to_dict())
-
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 795e8421..e584feb9 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -191,10 +191,13 @@ def _print_test_names(test_classes):
 
 
 def _record_suite_info(log_path, suite_class_name):
+  """Adds the record for suite information in test summary file."""
   summary_path = os.path.join(log_path, records.OUTPUT_FILE_SUMMARY)
   summary_writer = records.TestSummaryWriter(summary_path)
   content = records.SuiteInfoRecord(suite_class_name=suite_class_name)
-  summary_writer.dump(content.to_dict(), records.TestSummaryEntryType.SUITE_INFO)
+  summary_writer.dump(
+      content.to_dict(), records.TestSummaryEntryType.SUITE_INFO
+  )
 
 
 def run_suite_class(argv=None):
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index f9648bec..b3994819 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -542,7 +542,7 @@ def test_uid_helper():
 
     self.assertEqual(test_uid_helper.uid, 'some-uuid')
 
-  def test_create_suite_info_record_and_convert_to_dict(self):
+  def test_convert_suite_info_record_to_dict(self):
     suite_class_name = 'FakeTestSuite'
     record = records.SuiteInfoRecord(suite_class_name=suite_class_name)
 
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 99f8af7d..e176437f 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -176,7 +176,9 @@ def test_run_suite_class_records_suite_class_name(self, mock_time, _):
     expected_summary_entry = records.SuiteInfoRecord(
         suite_class_name='FakeTestSuite'
     ).to_dict()
-    expected_summary_entry['Type'] = records.TestSummaryEntryType.SUITE_INFO.value
+    expected_summary_entry['Type'] = (
+        records.TestSummaryEntryType.SUITE_INFO.value
+    )
 
     class FakeTestSuite(base_suite.BaseSuite):
 

From 3985e78b8aed6045177128b90171ca622970338d Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Tue, 3 Dec 2024 19:26:26 +0800
Subject: [PATCH 03/13] Fix comments. Add begin time and end time to suite info
 record.

---
 mobly/records.py                 | 31 +++++++++++++++++++++++--------
 mobly/suite_runner.py            | 12 +++++++-----
 tests/mobly/records_test.py      |  8 ++++++--
 tests/mobly/suite_runner_test.py | 12 ++++++++----
 4 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/mobly/records.py b/mobly/records.py
index 970e2ea6..59fd5a52 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -729,19 +729,34 @@ def summary_dict(self):
 
 
 class SuiteInfoRecord:
-  """A record representing the suite info in test summary."""
+  """A record representing the suite info in test summary.
 
-  KEY_SUITE_CLASS_NAME = 'Suite Class Name'
-  KEY_TIMESTAMP = 'Timestamp'
+  Attributes:
+    suite_class: str, the class name of the test suite class.
+    begin_time: int, epoch timestamp of when the suite started.
+    end_time: int, epoch timestamp of when the suite ended.
+  """
 
-  def __init__(self, suite_class_name):
-    self.suite_class_name = suite_class_name
-    self.timestamp = time.time()
+  KEY_SUITE_CLASS = 'Suite Class'
+  KEY_BEGIN_TIME = TestResultEnums.RECORD_BEGIN_TIME
+  KEY_END_TIME = TestResultEnums.RECORD_END_TIME
+
+  def __init__(self, suite_class):
+    self.suite_class = suite_class
+    self.begin_time = None
+    self.end_time = None
+
+  def suite_begin(self):
+    self.begin_time = utils.get_current_epoch_time()
+
+  def suite_end(self):
+    self.end_time = utils.get_current_epoch_time()
 
   def to_dict(self):
     result = {}
-    result[self.KEY_SUITE_CLASS_NAME] = self.suite_class_name
-    result[self.KEY_TIMESTAMP] = self.timestamp
+    result[self.KEY_SUITE_CLASS] = self.suite_class
+    result[self.KEY_BEGIN_TIME] = self.begin_time
+    result[self.KEY_END_TIME] = self.end_time
     return result
 
   def __repr__(self):
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index e584feb9..6e792f75 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -190,13 +190,12 @@ def _print_test_names(test_classes):
       print(f'{cls.TAG}.{name}')
 
 
-def _record_suite_info(log_path, suite_class_name):
-  """Adds the record for suite information in test summary file."""
+def _record_suite_info(suite_record, log_path):
+  """Dumps the suite info record to test summary file."""
   summary_path = os.path.join(log_path, records.OUTPUT_FILE_SUMMARY)
   summary_writer = records.TestSummaryWriter(summary_path)
-  content = records.SuiteInfoRecord(suite_class_name=suite_class_name)
   summary_writer.dump(
-      content.to_dict(), records.TestSummaryEntryType.SUITE_INFO
+      suite_record.to_dict(), records.TestSummaryEntryType.SUITE_INFO
   )
 
 
@@ -224,11 +223,12 @@ def run_suite_class(argv=None):
   suite = suite_class(runner, config)
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
+  suite_record = records.SuiteInfoRecord(suite_class=suite_class.__name__)
   with runner.mobly_logger(console_level=console_level) as log_path:
-    _record_suite_info(log_path, suite_class_name=suite_class.__name__)
     try:
       suite.setup_suite(config.copy())
       try:
+        suite_record.suite_begin()
         runner.run()
         ok = runner.results.is_all_pass
         print(ok)
@@ -236,6 +236,8 @@ def run_suite_class(argv=None):
         pass
     finally:
       suite.teardown_suite()
+      suite_record.suite_end()
+      _record_suite_info(suite_record, log_path)
   if not ok:
     sys.exit(1)
 
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index b3994819..65bfe05a 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -544,14 +544,18 @@ def test_uid_helper():
 
   def test_convert_suite_info_record_to_dict(self):
     suite_class_name = 'FakeTestSuite'
-    record = records.SuiteInfoRecord(suite_class_name=suite_class_name)
+    record = records.SuiteInfoRecord(suite_class=suite_class_name)
+    record.suite_begin()
+    record.suite_end()
 
     result = record.to_dict()
 
     self.assertIn(
-        (records.SuiteInfoRecord.KEY_SUITE_CLASS_NAME, suite_class_name),
+        (records.SuiteInfoRecord.KEY_SUITE_CLASS, suite_class_name),
         result.items(),
     )
+    self.assertIn(records.SuiteInfoRecord.KEY_BEGIN_TIME, result)
+    self.assertIn(records.SuiteInfoRecord.KEY_END_TIME, result)
 
 
 if __name__ == '__main__':
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index e176437f..49c8d0b1 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -26,6 +26,7 @@
 from mobly import base_test
 from mobly import records
 from mobly import suite_runner
+from mobly import utils
 from tests.lib import integration2_test
 from tests.lib import integration_test
 import yaml
@@ -169,13 +170,16 @@ def teardown_suite(self):
     mock_exit.assert_not_called()
 
   @mock.patch('sys.exit')
-  @mock.patch.object(time, 'time', return_value=1733143236.278318)
+  @mock.patch.object(
+      utils, 'get_current_epoch_time', return_value=1733143236278
+  )
   def test_run_suite_class_records_suite_class_name(self, mock_time, _):
     tmp_file_path = self._gen_tmp_config_file()
     mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
-    expected_summary_entry = records.SuiteInfoRecord(
-        suite_class_name='FakeTestSuite'
-    ).to_dict()
+    expected_record = records.SuiteInfoRecord(suite_class='FakeTestSuite')
+    expected_record.suite_begin()
+    expected_record.suite_end()
+    expected_summary_entry = expected_record.to_dict()
     expected_summary_entry['Type'] = (
         records.TestSummaryEntryType.SUITE_INFO.value
     )

From 74e4363d49bd5e6955e3993ce06579d8c892122e Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Tue, 3 Dec 2024 21:03:59 +0800
Subject: [PATCH 04/13] fix comments

---
 mobly/records.py                 | 20 ++++++++++++++------
 mobly/suite_runner.py            |  2 +-
 tests/mobly/records_test.py      |  4 ++--
 tests/mobly/suite_runner_test.py |  2 +-
 4 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/mobly/records.py b/mobly/records.py
index 59fd5a52..d0eef8ab 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -729,32 +729,40 @@ def summary_dict(self):
 
 
 class SuiteInfoRecord:
-  """A record representing the suite info in test summary.
+  """A record representing the test suite info in test summary.
 
   Attributes:
-    suite_class: str, the class name of the test suite class.
+    test_suite_class: str, the class name of the test suite class.
     begin_time: int, epoch timestamp of when the suite started.
     end_time: int, epoch timestamp of when the suite ended.
   """
 
-  KEY_SUITE_CLASS = 'Suite Class'
+  KEY_TEST_SUITE_CLASS = 'Test Suite Class'
   KEY_BEGIN_TIME = TestResultEnums.RECORD_BEGIN_TIME
   KEY_END_TIME = TestResultEnums.RECORD_END_TIME
 
-  def __init__(self, suite_class):
-    self.suite_class = suite_class
+  def __init__(self, test_suite_class):
+    self.test_suite_class = test_suite_class
     self.begin_time = None
     self.end_time = None
 
   def suite_begin(self):
+    """Call this when the suite begins execution.
+
+    Sets the begin_time of this record.
+    """
     self.begin_time = utils.get_current_epoch_time()
 
   def suite_end(self):
+    """Call this when the suite ends execution.
+
+    Sets the end_time of this record.
+    """
     self.end_time = utils.get_current_epoch_time()
 
   def to_dict(self):
     result = {}
-    result[self.KEY_SUITE_CLASS] = self.suite_class
+    result[self.KEY_TEST_SUITE_CLASS] = self.test_suite_class
     result[self.KEY_BEGIN_TIME] = self.begin_time
     result[self.KEY_END_TIME] = self.end_time
     return result
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 6e792f75..03818422 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -223,7 +223,7 @@ def run_suite_class(argv=None):
   suite = suite_class(runner, config)
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
-  suite_record = records.SuiteInfoRecord(suite_class=suite_class.__name__)
+  suite_record = records.SuiteInfoRecord(test_suite_class=suite_class.__name__)
   with runner.mobly_logger(console_level=console_level) as log_path:
     try:
       suite.setup_suite(config.copy())
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index 65bfe05a..5ebb7ecf 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -544,14 +544,14 @@ def test_uid_helper():
 
   def test_convert_suite_info_record_to_dict(self):
     suite_class_name = 'FakeTestSuite'
-    record = records.SuiteInfoRecord(suite_class=suite_class_name)
+    record = records.SuiteInfoRecord(test_suite_class=suite_class_name)
     record.suite_begin()
     record.suite_end()
 
     result = record.to_dict()
 
     self.assertIn(
-        (records.SuiteInfoRecord.KEY_SUITE_CLASS, suite_class_name),
+        (records.SuiteInfoRecord.KEY_TEST_SUITE_CLASS, suite_class_name),
         result.items(),
     )
     self.assertIn(records.SuiteInfoRecord.KEY_BEGIN_TIME, result)
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 49c8d0b1..3e72691d 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -176,7 +176,7 @@ def teardown_suite(self):
   def test_run_suite_class_records_suite_class_name(self, mock_time, _):
     tmp_file_path = self._gen_tmp_config_file()
     mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
-    expected_record = records.SuiteInfoRecord(suite_class='FakeTestSuite')
+    expected_record = records.SuiteInfoRecord(test_suite_class='FakeTestSuite')
     expected_record.suite_begin()
     expected_record.suite_end()
     expected_summary_entry = expected_record.to_dict()

From efb5d762157120cfbc80fbe700108046fb25fc9f Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Wed, 4 Dec 2024 16:43:41 +0800
Subject: [PATCH 05/13] Move SuiteInfoRecord class to suite_runner; Add a
 VERSION attribute to BaseSuite and record it to test summary

---
 mobly/base_suite.py              |  6 ++++
 mobly/records.py                 | 45 -----------------------
 mobly/suite_runner.py            | 61 +++++++++++++++++++++++++++++---
 tests/mobly/records_test.py      | 15 --------
 tests/mobly/suite_runner_test.py | 28 +++++++++++++--
 5 files changed, 89 insertions(+), 66 deletions(-)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index f4399ac7..609b1cc4 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -29,8 +29,14 @@ class BaseSuite(abc.ABC):
   Users can use this class if they need to define their own setup and teardown
   steps on the suite level. Otherwise, just use suite_runner.run_suite on the
   list of test classes.
+
+  Attributes:
+    VERSION: str, the test suite version string. Default to None. Subclasses may
+      overwrite this attribute to record suite version in test summary.
   """
 
+  VERSION = None
+
   def __init__(self, runner, config):
     self._runner = runner
     self._config = config.copy()
diff --git a/mobly/records.py b/mobly/records.py
index d0eef8ab..10d3063c 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -94,8 +94,6 @@ class TestSummaryEntryType(enum.Enum):
   SUMMARY = 'Summary'
   # Information on the controllers used in a test class.
   CONTROLLER_INFO = 'ControllerInfo'
-  # Test suite level information.
-  SUITE_INFO = 'SuiteInfo'
   # Additional data added by users during test.
   # This can be added at any point in the test, so do not assume the location
   # of these entries in the summary file.
@@ -726,46 +724,3 @@ def summary_dict(self):
     d['Skipped'] = len(self.skipped)
     d['Error'] = len(self.error)
     return d
-
-
-class SuiteInfoRecord:
-  """A record representing the test suite info in test summary.
-
-  Attributes:
-    test_suite_class: str, the class name of the test suite class.
-    begin_time: int, epoch timestamp of when the suite started.
-    end_time: int, epoch timestamp of when the suite ended.
-  """
-
-  KEY_TEST_SUITE_CLASS = 'Test Suite Class'
-  KEY_BEGIN_TIME = TestResultEnums.RECORD_BEGIN_TIME
-  KEY_END_TIME = TestResultEnums.RECORD_END_TIME
-
-  def __init__(self, test_suite_class):
-    self.test_suite_class = test_suite_class
-    self.begin_time = None
-    self.end_time = None
-
-  def suite_begin(self):
-    """Call this when the suite begins execution.
-
-    Sets the begin_time of this record.
-    """
-    self.begin_time = utils.get_current_epoch_time()
-
-  def suite_end(self):
-    """Call this when the suite ends execution.
-
-    Sets the end_time of this record.
-    """
-    self.end_time = utils.get_current_epoch_time()
-
-  def to_dict(self):
-    result = {}
-    result[self.KEY_TEST_SUITE_CLASS] = self.test_suite_class
-    result[self.KEY_BEGIN_TIME] = self.begin_time
-    result[self.KEY_END_TIME] = self.end_time
-    return result
-
-  def __repr__(self):
-    return str(self.to_dict())
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 03818422..2c5d17f3 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -77,12 +77,60 @@ def setup_suite(self, config):
 from mobly import records
 from mobly import signals
 from mobly import test_runner
+from mobly import utils
 
 
 class Error(Exception):
   pass
 
 
+class SuiteInfoRecord:
+  """A record representing the test suite info in test summary.
+
+  Attributes:
+    test_suite_class: str, the class name of the test suite class.
+    test_suite_version: str, the test suite version.
+    begin_time: int, epoch timestamp of when the suite started.
+    end_time: int, epoch timestamp of when the suite ended.
+  """
+
+  KEY_TEST_SUITE_CLASS = 'Test Suite Class'
+  KEY_TEST_SUITE_VERSION = 'Test Suite Version'
+  KEY_BEGIN_TIME = records.TestResultEnums.RECORD_BEGIN_TIME
+  KEY_END_TIME = records.TestResultEnums.RECORD_END_TIME
+
+  def __init__(self, test_suite_class, test_suite_version=None):
+    self.test_suite_class = test_suite_class
+    self.test_suite_version = test_suite_version
+    self.begin_time = None
+    self.end_time = None
+
+  def suite_begin(self):
+    """Call this when the suite begins execution.
+
+    Sets the begin_time of this record.
+    """
+    self.begin_time = utils.get_current_epoch_time()
+
+  def suite_end(self):
+    """Call this when the suite ends execution.
+
+    Sets the end_time of this record.
+    """
+    self.end_time = utils.get_current_epoch_time()
+
+  def to_dict(self):
+    result = {}
+    result[self.KEY_TEST_SUITE_CLASS] = self.test_suite_class
+    result[self.KEY_TEST_SUITE_VERSION] = self.test_suite_version
+    result[self.KEY_BEGIN_TIME] = self.begin_time
+    result[self.KEY_END_TIME] = self.end_time
+    return result
+
+  def __repr__(self):
+    return str(self.to_dict())
+
+
 def _parse_cli_args(argv):
   """Parses cli args that are consumed by Mobly.
 
@@ -190,12 +238,12 @@ def _print_test_names(test_classes):
       print(f'{cls.TAG}.{name}')
 
 
-def _record_suite_info(suite_record, log_path):
+def _dump_suite_info(suite_record, log_path):
   """Dumps the suite info record to test summary file."""
   summary_path = os.path.join(log_path, records.OUTPUT_FILE_SUMMARY)
   summary_writer = records.TestSummaryWriter(summary_path)
   summary_writer.dump(
-      suite_record.to_dict(), records.TestSummaryEntryType.SUITE_INFO
+      suite_record.to_dict(), records.TestSummaryEntryType.USER_DATA
   )
 
 
@@ -221,9 +269,14 @@ def run_suite_class(argv=None):
       log_dir=config.log_path, testbed_name=config.testbed_name
   )
   suite = suite_class(runner, config)
+
+  suite_version = getattr(suite_class, 'VERSION', None)
+  suite_record = SuiteInfoRecord(
+      test_suite_class=suite_class.__name__, test_suite_version=suite_version
+  )
+
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
-  suite_record = records.SuiteInfoRecord(test_suite_class=suite_class.__name__)
   with runner.mobly_logger(console_level=console_level) as log_path:
     try:
       suite.setup_suite(config.copy())
@@ -237,7 +290,7 @@ def run_suite_class(argv=None):
     finally:
       suite.teardown_suite()
       suite_record.suite_end()
-      _record_suite_info(suite_record, log_path)
+      _dump_suite_info(suite_record, log_path)
   if not ok:
     sys.exit(1)
 
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index 5ebb7ecf..5c1d027f 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -542,21 +542,6 @@ def test_uid_helper():
 
     self.assertEqual(test_uid_helper.uid, 'some-uuid')
 
-  def test_convert_suite_info_record_to_dict(self):
-    suite_class_name = 'FakeTestSuite'
-    record = records.SuiteInfoRecord(test_suite_class=suite_class_name)
-    record.suite_begin()
-    record.suite_end()
-
-    result = record.to_dict()
-
-    self.assertIn(
-        (records.SuiteInfoRecord.KEY_TEST_SUITE_CLASS, suite_class_name),
-        result.items(),
-    )
-    self.assertIn(records.SuiteInfoRecord.KEY_BEGIN_TIME, result)
-    self.assertIn(records.SuiteInfoRecord.KEY_END_TIME, result)
-
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 3e72691d..4ddb035a 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -176,12 +176,14 @@ def teardown_suite(self):
   def test_run_suite_class_records_suite_class_name(self, mock_time, _):
     tmp_file_path = self._gen_tmp_config_file()
     mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
-    expected_record = records.SuiteInfoRecord(test_suite_class='FakeTestSuite')
+    expected_record = suite_runner.SuiteInfoRecord(
+        test_suite_class='FakeTestSuite'
+    )
     expected_record.suite_begin()
     expected_record.suite_end()
     expected_summary_entry = expected_record.to_dict()
     expected_summary_entry['Type'] = (
-        records.TestSummaryEntryType.SUITE_INFO.value
+        records.TestSummaryEntryType.USER_DATA.value
     )
 
     class FakeTestSuite(base_suite.BaseSuite):
@@ -225,6 +227,28 @@ def test_print_test_names_with_exception(self):
     mock_cls_instance._pre_run.side_effect = Exception('Something went wrong.')
     mock_cls_instance._clean_up.assert_called_once()
 
+  def test_convert_suite_info_record_to_dict(self):
+    suite_class_name = 'FakeTestSuite'
+    suite_version = '1.2.3'
+    record = suite_runner.SuiteInfoRecord(
+        test_suite_class=suite_class_name, test_suite_version=suite_version
+    )
+    record.suite_begin()
+    record.suite_end()
+
+    result = record.to_dict()
+
+    self.assertIn(
+        (suite_runner.SuiteInfoRecord.KEY_TEST_SUITE_CLASS, suite_class_name),
+        result.items(),
+    )
+    self.assertIn(
+        (suite_runner.SuiteInfoRecord.KEY_TEST_SUITE_VERSION, suite_version),
+        result.items(),
+    )
+    self.assertIn(suite_runner.SuiteInfoRecord.KEY_BEGIN_TIME, result)
+    self.assertIn(suite_runner.SuiteInfoRecord.KEY_END_TIME, result)
+
   def _gen_tmp_config_file(self):
     tmp_file_path = os.path.join(self.tmp_dir, 'config.yml')
     with io.open(tmp_file_path, 'w', encoding='utf-8') as f:

From c272a9a1e23cb8e009e37b3226253af6e8b66660 Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Wed, 4 Dec 2024 20:45:44 +0800
Subject: [PATCH 06/13] Add a property BaseSuite.suite_info to enable dumping
 user defined suite info.

---
 mobly/base_suite.py              | 11 +++----
 mobly/suite_runner.py            | 55 +++++++++++++++-----------------
 tests/mobly/suite_runner_test.py |  4 +--
 3 files changed, 32 insertions(+), 38 deletions(-)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index 609b1cc4..6135804f 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -29,18 +29,17 @@ class BaseSuite(abc.ABC):
   Users can use this class if they need to define their own setup and teardown
   steps on the suite level. Otherwise, just use suite_runner.run_suite on the
   list of test classes.
-
-  Attributes:
-    VERSION: str, the test suite version string. Default to None. Subclasses may
-      overwrite this attribute to record suite version in test summary.
   """
 
-  VERSION = None
-
   def __init__(self, runner, config):
     self._runner = runner
     self._config = config.copy()
 
+  @property
+  def suite_info(self):
+    """Returns user defined suite info that will be recorded to test summary."""
+    return {}
+
   @property
   def user_params(self):
     return self._config.user_params
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 2c5d17f3..7dcceefb 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -85,46 +85,42 @@ class Error(Exception):
 
 
 class SuiteInfoRecord:
-  """A record representing the test suite info in test summary.
-
-  Attributes:
-    test_suite_class: str, the class name of the test suite class.
-    test_suite_version: str, the test suite version.
-    begin_time: int, epoch timestamp of when the suite started.
-    end_time: int, epoch timestamp of when the suite ended.
-  """
+  """A record representing the test suite info in test summary."""
 
   KEY_TEST_SUITE_CLASS = 'Test Suite Class'
-  KEY_TEST_SUITE_VERSION = 'Test Suite Version'
+  KEY_EXTRAS = records.TestResultEnums.RECORD_EXTRAS
   KEY_BEGIN_TIME = records.TestResultEnums.RECORD_BEGIN_TIME
   KEY_END_TIME = records.TestResultEnums.RECORD_END_TIME
 
-  def __init__(self, test_suite_class, test_suite_version=None):
-    self.test_suite_class = test_suite_class
-    self.test_suite_version = test_suite_version
-    self.begin_time = None
-    self.end_time = None
+  # The class name of the test suite class.
+  _test_suite_class: str
+  # Epoch timestamp of when the suite started.
+  _begin_time: int
+  # Epoch timestamp of when the suite ended.
+  _end_time: int
+  # User defined extra information of the test result. Must be serializable.
+  _extras: dict
+
+  def __init__(self, test_suite_class, extras=None):
+    self._test_suite_class = test_suite_class
+    self._extras = extras or dict()
+    self._begin_time = None
+    self._end_time = None
 
   def suite_begin(self):
-    """Call this when the suite begins execution.
-
-    Sets the begin_time of this record.
-    """
-    self.begin_time = utils.get_current_epoch_time()
+    """Call this when the suite begins execution."""
+    self._begin_time = utils.get_current_epoch_time()
 
   def suite_end(self):
-    """Call this when the suite ends execution.
-
-    Sets the end_time of this record.
-    """
-    self.end_time = utils.get_current_epoch_time()
+    """Call this when the suite ends execution."""
+    self._end_time = utils.get_current_epoch_time()
 
   def to_dict(self):
     result = {}
-    result[self.KEY_TEST_SUITE_CLASS] = self.test_suite_class
-    result[self.KEY_TEST_SUITE_VERSION] = self.test_suite_version
-    result[self.KEY_BEGIN_TIME] = self.begin_time
-    result[self.KEY_END_TIME] = self.end_time
+    result[self.KEY_TEST_SUITE_CLASS] = self._test_suite_class
+    result[self.KEY_EXTRAS] = self._extras
+    result[self.KEY_BEGIN_TIME] = self._begin_time
+    result[self.KEY_END_TIME] = self._end_time
     return result
 
   def __repr__(self):
@@ -270,9 +266,8 @@ def run_suite_class(argv=None):
   )
   suite = suite_class(runner, config)
 
-  suite_version = getattr(suite_class, 'VERSION', None)
   suite_record = SuiteInfoRecord(
-      test_suite_class=suite_class.__name__, test_suite_version=suite_version
+      test_suite_class=suite_class.__name__, extras=suite.suite_info
   )
 
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 4ddb035a..23663067 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -231,7 +231,7 @@ def test_convert_suite_info_record_to_dict(self):
     suite_class_name = 'FakeTestSuite'
     suite_version = '1.2.3'
     record = suite_runner.SuiteInfoRecord(
-        test_suite_class=suite_class_name, test_suite_version=suite_version
+        test_suite_class=suite_class_name, extras={'version': suite_version},
     )
     record.suite_begin()
     record.suite_end()
@@ -243,7 +243,7 @@ def test_convert_suite_info_record_to_dict(self):
         result.items(),
     )
     self.assertIn(
-        (suite_runner.SuiteInfoRecord.KEY_TEST_SUITE_VERSION, suite_version),
+        (suite_runner.SuiteInfoRecord.KEY_EXTRAS, {'version': suite_version}),
         result.items(),
     )
     self.assertIn(suite_runner.SuiteInfoRecord.KEY_BEGIN_TIME, result)

From 34c1787c0be93615ba843684a79b6ea192283b9c Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Wed, 4 Dec 2024 20:48:14 +0800
Subject: [PATCH 07/13] Fix format.

---
 tests/mobly/suite_runner_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 23663067..1baa8cd6 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -231,7 +231,7 @@ def test_convert_suite_info_record_to_dict(self):
     suite_class_name = 'FakeTestSuite'
     suite_version = '1.2.3'
     record = suite_runner.SuiteInfoRecord(
-        test_suite_class=suite_class_name, extras={'version': suite_version},
+        test_suite_class=suite_class_name, extras={'version': suite_version}
     )
     record.suite_begin()
     record.suite_end()

From 50003aa4fe122e3e3c0e38ac15f2db35efe08b6a Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Thu, 12 Dec 2024 16:45:15 +0800
Subject: [PATCH 08/13] Use a dedicate test summary type for suite info record

---
 mobly/suite_runner.py            | 15 +++++++++++----
 tests/mobly/suite_runner_test.py |  2 +-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 7dcceefb..6010febb 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -66,6 +66,7 @@ def setup_suite(self, config):
 """
 import argparse
 import collections
+import enum
 import inspect
 import logging
 import os
@@ -84,13 +85,19 @@ class Error(Exception):
   pass
 
 
+class TestSummaryEntryType(enum.Enum):
+  """Constants used to record suite level entries in test summary file."""
+
+  SUITE_INFO = 'SuiteInfo'
+
+
 class SuiteInfoRecord:
   """A record representing the test suite info in test summary."""
 
   KEY_TEST_SUITE_CLASS = 'Test Suite Class'
-  KEY_EXTRAS = records.TestResultEnums.RECORD_EXTRAS
-  KEY_BEGIN_TIME = records.TestResultEnums.RECORD_BEGIN_TIME
-  KEY_END_TIME = records.TestResultEnums.RECORD_END_TIME
+  KEY_EXTRAS = 'Extras'
+  KEY_BEGIN_TIME = 'Suite Begin Time'
+  KEY_END_TIME = 'Suite End Time'
 
   # The class name of the test suite class.
   _test_suite_class: str
@@ -239,7 +246,7 @@ def _dump_suite_info(suite_record, log_path):
   summary_path = os.path.join(log_path, records.OUTPUT_FILE_SUMMARY)
   summary_writer = records.TestSummaryWriter(summary_path)
   summary_writer.dump(
-      suite_record.to_dict(), records.TestSummaryEntryType.USER_DATA
+      suite_record.to_dict(), TestSummaryEntryType.SUITE_INFO
   )
 
 
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 1baa8cd6..5035d640 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -183,7 +183,7 @@ def test_run_suite_class_records_suite_class_name(self, mock_time, _):
     expected_record.suite_end()
     expected_summary_entry = expected_record.to_dict()
     expected_summary_entry['Type'] = (
-        records.TestSummaryEntryType.USER_DATA.value
+        suite_runner.TestSummaryEntryType.SUITE_INFO.value
     )
 
     class FakeTestSuite(base_suite.BaseSuite):

From 068482893b439621013c2849beeef1d64c18d801 Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Thu, 12 Dec 2024 16:47:35 +0800
Subject: [PATCH 09/13] fix format

---
 mobly/suite_runner.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 6010febb..178a91b8 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -245,9 +245,7 @@ def _dump_suite_info(suite_record, log_path):
   """Dumps the suite info record to test summary file."""
   summary_path = os.path.join(log_path, records.OUTPUT_FILE_SUMMARY)
   summary_writer = records.TestSummaryWriter(summary_path)
-  summary_writer.dump(
-      suite_record.to_dict(), TestSummaryEntryType.SUITE_INFO
-  )
+  summary_writer.dump(suite_record.to_dict(), TestSummaryEntryType.SUITE_INFO)
 
 
 def run_suite_class(argv=None):

From f9bce18f312d15d4be25baae096b0c18719f717d Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Wed, 18 Dec 2024 00:16:46 +0800
Subject: [PATCH 10/13] Modify API name and docstring

---
 mobly/base_suite.py   | 15 ++++++++++-----
 mobly/suite_runner.py | 12 +++++++-----
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index 6135804f..188a9e0d 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -35,15 +35,20 @@ def __init__(self, runner, config):
     self._runner = runner
     self._config = config.copy()
 
-  @property
-  def suite_info(self):
-    """Returns user defined suite info that will be recorded to test summary."""
-    return {}
-
   @property
   def user_params(self):
     return self._config.user_params
 
+  def get_suite_info(self):
+    """User defined extra suite information to be recorded in test summary.
+
+    This method will be called after all test classes are executed.
+
+    Returns:
+      A dict of suite information. Keys and values must be serializable.
+    """
+    return {}
+
   def add_test_class(self, clazz, config=None, tests=None, name_suffix=None):
     """Adds a test class to the suite.
 
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 178a91b8..9a1f0e2c 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -108,9 +108,9 @@ class SuiteInfoRecord:
   # User defined extra information of the test result. Must be serializable.
   _extras: dict
 
-  def __init__(self, test_suite_class, extras=None):
+  def __init__(self, test_suite_class):
     self._test_suite_class = test_suite_class
-    self._extras = extras or dict()
+    self._extras = dict()
     self._begin_time = None
     self._end_time = None
 
@@ -122,6 +122,9 @@ def suite_end(self):
     """Call this when the suite ends execution."""
     self._end_time = utils.get_current_epoch_time()
 
+  def set_extras(self, extras):
+    self._extras = extras
+
   def to_dict(self):
     result = {}
     result[self.KEY_TEST_SUITE_CLASS] = self._test_suite_class
@@ -271,9 +274,7 @@ def run_suite_class(argv=None):
   )
   suite = suite_class(runner, config)
 
-  suite_record = SuiteInfoRecord(
-      test_suite_class=suite_class.__name__, extras=suite.suite_info
-  )
+  suite_record = SuiteInfoRecord(test_suite_class=suite_class.__name__)
 
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
@@ -290,6 +291,7 @@ def run_suite_class(argv=None):
     finally:
       suite.teardown_suite()
       suite_record.suite_end()
+      suite_record.set_extras(suite.get_suite_info())
       _dump_suite_info(suite_record, log_path)
   if not ok:
     sys.exit(1)

From 929cfe808b48aca4ebacad86b4cae8d9323a4338 Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Wed, 18 Dec 2024 01:06:25 +0800
Subject: [PATCH 11/13] support customized suite name

---
 mobly/base_suite.py              | 12 ++++++++++++
 mobly/suite_runner.py            | 19 +++++++++++++++----
 tests/mobly/suite_runner_test.py | 18 +++++++++++++++++-
 3 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index 32442082..9808b32f 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -50,6 +50,18 @@ def set_test_selector(self, test_selector):
     """
     self._test_selector = test_selector
 
+  def get_suite_name(self):
+    """The name of this suite.
+
+    By default, use the name of the test class. User can overwrite to return
+    a customized suite name. User can include test runtime info as this will be
+    collected after all test classes are executed.
+
+    Returns:
+      A string of suite name.
+    """
+    return self.__class__.__name__
+
   def get_suite_info(self):
     """User defined extra suite information to be recorded in test summary.
 
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 8ff90bf1..0ba3b04b 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -94,22 +94,26 @@ class TestSummaryEntryType(enum.Enum):
 class SuiteInfoRecord:
   """A record representing the test suite info in test summary."""
 
+  KEY_SUITE_NAME = 'Suite Name'
   KEY_TEST_SUITE_CLASS = 'Test Suite Class'
   KEY_EXTRAS = 'Extras'
   KEY_BEGIN_TIME = 'Suite Begin Time'
   KEY_END_TIME = 'Suite End Time'
 
+  # The name of the test suite.
+  _suite_name: str
   # The class name of the test suite class.
   _test_suite_class: str
-  # Epoch timestamp of when the suite started.
-  _begin_time: int
-  # Epoch timestamp of when the suite ended.
-  _end_time: int
   # User defined extra information of the test result. Must be serializable.
   _extras: dict
+  # Epoch timestamp of when the suite started.
+  _begin_time: int | None
+  # Epoch timestamp of when the suite ended.
+  _end_time: int | None
 
   def __init__(self, test_suite_class):
     self._test_suite_class = test_suite_class
+    self._suite_name = ''
     self._extras = dict()
     self._begin_time = None
     self._end_time = None
@@ -122,12 +126,18 @@ def suite_end(self):
     """Call this when the suite ends execution."""
     self._end_time = utils.get_current_epoch_time()
 
+  def set_suite_name(self, suite_name):
+    """Sets the name of the test suite."""
+    self._suite_name = suite_name
+
   def set_extras(self, extras):
+    """Sets extra information. Must be serializable."""
     self._extras = extras
 
   def to_dict(self):
     result = {}
     result[self.KEY_TEST_SUITE_CLASS] = self._test_suite_class
+    result[self.KEY_SUITE_NAME] = self._suite_name
     result[self.KEY_EXTRAS] = self._extras
     result[self.KEY_BEGIN_TIME] = self._begin_time
     result[self.KEY_END_TIME] = self._end_time
@@ -334,6 +344,7 @@ def run_suite_class(argv=None):
     finally:
       suite.teardown_suite()
       suite_record.suite_end()
+      suite_record.set_suite_name(suite.get_suite_name())
       suite_record.set_extras(suite.get_suite_info())
       _dump_suite_info(suite_record, log_path)
   if not ok:
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index deb0b20d..368ec0bf 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -260,7 +260,6 @@ def test_run_suite_class_finds_suite_class_when_not_in_main_module(
   ):
     mock_test_runner = mock_test_runner_class.return_value
     mock_test_runner.results.is_all_pass = True
-    mock_test_runner
     tmp_file_path = self._gen_tmp_config_file()
     mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
 
@@ -275,12 +274,20 @@ def test_run_suite_class_finds_suite_class_when_not_in_main_module(
   )
   def test_run_suite_class_records_suite_info(self, mock_time, _):
     tmp_file_path = self._gen_tmp_config_file()
+    customized_suite_name = 'Customized Suite Name'
     mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
     expected_record = suite_runner.SuiteInfoRecord(
         test_suite_class='FakeTestSuite'
     )
     expected_record.suite_begin()
     expected_record.suite_end()
+    expected_record.set_suite_name(customized_suite_name)
+    expected_record.set_extras(
+        {
+            'extra-key-0': 'extra-value-0',
+            'extra-key-1': 'extra-value-1',
+        }
+    )
     expected_summary_entry = expected_record.to_dict()
     expected_summary_entry['Type'] = (
         suite_runner.TestSummaryEntryType.SUITE_INFO.value
@@ -288,6 +295,15 @@ def test_run_suite_class_records_suite_info(self, mock_time, _):
 
     class FakeTestSuite(base_suite.BaseSuite):
 
+      def get_suite_name(self):
+        return customized_suite_name
+
+      def get_suite_info(self):
+        return {
+            'extra-key-0': 'extra-value-0',
+            'extra-key-1': 'extra-value-1',
+        }
+
       def setup_suite(self, config):
         super().setup_suite(config)
         self.add_test_class(FakeTest1)

From 4b432edc55536ffeea7602b04a9cc98899e29fe4 Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Tue, 24 Dec 2024 16:25:20 +0800
Subject: [PATCH 12/13] Fix comments; Add an interface get_run_identifier to
 BaseSuite class.

---
 mobly/base_suite.py              | 56 +++++++++++++++++++-------------
 mobly/suite_runner.py            | 34 ++++++++++++-------
 tests/mobly/suite_runner_test.py | 25 +++++++++++---
 3 files changed, 77 insertions(+), 38 deletions(-)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index 9808b32f..f2a7378a 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -50,28 +50,6 @@ def set_test_selector(self, test_selector):
     """
     self._test_selector = test_selector
 
-  def get_suite_name(self):
-    """The name of this suite.
-
-    By default, use the name of the test class. User can overwrite to return
-    a customized suite name. User can include test runtime info as this will be
-    collected after all test classes are executed.
-
-    Returns:
-      A string of suite name.
-    """
-    return self.__class__.__name__
-
-  def get_suite_info(self):
-    """User defined extra suite information to be recorded in test summary.
-
-    This method will be called after all test classes are executed.
-
-    Returns:
-      A dict of suite information. Keys and values must be serializable.
-    """
-    return {}
-
   def add_test_class(self, clazz, config=None, tests=None, name_suffix=None):
     """Adds a test class to the suite.
 
@@ -117,3 +95,37 @@ def setup_suite(self, config):
   def teardown_suite(self):
     """Function used to add post tests cleanup tasks (optional)."""
     pass
+
+  # Optional interfaces to record user defined suite information to test summary
+
+  def get_suite_name(self):
+    """Override to return a customized suite name (optional).
+
+    Use suite class name by default.
+
+    Returns:
+      A string that indicates the suite name.
+    """
+    return self.__class__.__name__
+
+  def get_run_identifier(self):
+    """Override to record identifier describing the key run context (optional).
+
+    Users can include test runtime info as this method will be called after all
+    test classes are executed.
+
+    Returns:
+      A string that indicates key run context information.
+    """
+    return None
+
+  def get_suite_info(self):
+    """Override to record user defined extra info to test summary (optional).
+
+    Users can include test runtime info as this method will be called after all
+    test classes are executed.
+
+    Returns:
+      A dict of suite information. Keys and values must be serializable.
+    """
+    return {}
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 0ba3b04b..67a7f13b 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -92,10 +92,15 @@ class TestSummaryEntryType(enum.Enum):
 
 
 class SuiteInfoRecord:
-  """A record representing the test suite info in test summary."""
+  """A record representing the test suite info in test summary.
+
+  This record class is for suites defined by inheriting `base_suite.BaseSuite`.
+  This is not for suites directly assembled via `run_suite`.
+  """
 
   KEY_SUITE_NAME = 'Suite Name'
-  KEY_TEST_SUITE_CLASS = 'Test Suite Class'
+  KEY_SUITE_CLASS_NAME = 'Suite Class Name'
+  KEY_RUN_IDENTIFIER = 'Run Identifier'
   KEY_EXTRAS = 'Extras'
   KEY_BEGIN_TIME = 'Suite Begin Time'
   KEY_END_TIME = 'Suite End Time'
@@ -103,20 +108,20 @@ class SuiteInfoRecord:
   # The name of the test suite.
   _suite_name: str
   # The class name of the test suite class.
-  _test_suite_class: str
+  _suite_class_name: str
+  # The run identifier that describes key test run context.
+  _run_identifier: str | None = None
   # User defined extra information of the test result. Must be serializable.
   _extras: dict
   # Epoch timestamp of when the suite started.
-  _begin_time: int | None
+  _begin_time: int | None = None
   # Epoch timestamp of when the suite ended.
-  _end_time: int | None
+  _end_time: int | None = None
 
-  def __init__(self, test_suite_class):
-    self._test_suite_class = test_suite_class
+  def __init__(self, suite_class_name):
+    self._suite_class_name = suite_class_name
     self._suite_name = ''
     self._extras = dict()
-    self._begin_time = None
-    self._end_time = None
 
   def suite_begin(self):
     """Call this when the suite begins execution."""
@@ -130,14 +135,18 @@ def set_suite_name(self, suite_name):
     """Sets the name of the test suite."""
     self._suite_name = suite_name
 
+  def set_run_identifier(self, run_identifier):
+    self._run_identifier = run_identifier
+
   def set_extras(self, extras):
     """Sets extra information. Must be serializable."""
     self._extras = extras
 
   def to_dict(self):
     result = {}
-    result[self.KEY_TEST_SUITE_CLASS] = self._test_suite_class
+    result[self.KEY_SUITE_CLASS_NAME] = self._suite_class_name
     result[self.KEY_SUITE_NAME] = self._suite_name
+    result[self.KEY_RUN_IDENTIFIER] = self._run_identifier
     result[self.KEY_EXTRAS] = self._extras
     result[self.KEY_BEGIN_TIME] = self._begin_time
     result[self.KEY_END_TIME] = self._end_time
@@ -327,7 +336,8 @@ def run_suite_class(argv=None):
   suite = suite_class(runner, config)
   test_selector = _parse_raw_test_selector(cli_args.tests)
   suite.set_test_selector(test_selector)
-  suite_record = SuiteInfoRecord(test_suite_class=suite_class.__name__)
+  suite_record = SuiteInfoRecord(suite_class_name=suite_class.__name__)
+  suite_record.set_suite_name(suite.get_suite_name())
 
   console_level = logging.DEBUG if cli_args.verbose else logging.INFO
   ok = False
@@ -344,7 +354,7 @@ def run_suite_class(argv=None):
     finally:
       suite.teardown_suite()
       suite_record.suite_end()
-      suite_record.set_suite_name(suite.get_suite_name())
+      suite_record.set_run_identifier(suite.get_run_identifier())
       suite_record.set_extras(suite.get_suite_info())
       _dump_suite_info(suite_record, log_path)
   if not ok:
diff --git a/tests/mobly/suite_runner_test.py b/tests/mobly/suite_runner_test.py
index 368ec0bf..d10337a2 100755
--- a/tests/mobly/suite_runner_test.py
+++ b/tests/mobly/suite_runner_test.py
@@ -275,13 +275,15 @@ def test_run_suite_class_finds_suite_class_when_not_in_main_module(
   def test_run_suite_class_records_suite_info(self, mock_time, _):
     tmp_file_path = self._gen_tmp_config_file()
     customized_suite_name = 'Customized Suite Name'
+    run_identifier = '123456'
     mock_cli_args = ['test_binary', f'--config={tmp_file_path}']
     expected_record = suite_runner.SuiteInfoRecord(
-        test_suite_class='FakeTestSuite'
+        suite_class_name='FakeTestSuite'
     )
+    expected_record.set_suite_name(customized_suite_name)
     expected_record.suite_begin()
     expected_record.suite_end()
-    expected_record.set_suite_name(customized_suite_name)
+    expected_record.set_run_identifier(run_identifier)
     expected_record.set_extras(
         {
             'extra-key-0': 'extra-value-0',
@@ -298,6 +300,9 @@ class FakeTestSuite(base_suite.BaseSuite):
       def get_suite_name(self):
         return customized_suite_name
 
+      def get_run_identifier(self):
+        return run_identifier
+
       def get_suite_info(self):
         return {
             'extra-key-0': 'extra-value-0',
@@ -345,22 +350,34 @@ def test_print_test_names_with_exception(self):
 
   def test_convert_suite_info_record_to_dict(self):
     suite_class_name = 'FakeTestSuite'
+    suite_name = 'Customized Suite Name'
+    run_identifier = '123456'
     suite_version = '1.2.3'
-    record = suite_runner.SuiteInfoRecord(test_suite_class=suite_class_name)
+    record = suite_runner.SuiteInfoRecord(suite_class_name=suite_class_name)
     record.set_extras({'version': suite_version})
+    record.set_suite_name(suite_name)
     record.suite_begin()
     record.suite_end()
+    record.set_run_identifier(run_identifier)
 
     result = record.to_dict()
 
     self.assertIn(
-        (suite_runner.SuiteInfoRecord.KEY_TEST_SUITE_CLASS, suite_class_name),
+        (suite_runner.SuiteInfoRecord.KEY_SUITE_CLASS_NAME, suite_class_name),
         result.items(),
     )
     self.assertIn(
         (suite_runner.SuiteInfoRecord.KEY_EXTRAS, {'version': suite_version}),
         result.items(),
     )
+    self.assertIn(
+        (suite_runner.SuiteInfoRecord.KEY_SUITE_NAME, suite_name),
+        result.items(),
+    )
+    self.assertIn(
+        (suite_runner.SuiteInfoRecord.KEY_RUN_IDENTIFIER, run_identifier),
+        result.items(),
+    )
     self.assertIn(suite_runner.SuiteInfoRecord.KEY_BEGIN_TIME, result)
     self.assertIn(suite_runner.SuiteInfoRecord.KEY_END_TIME, result)
 

From 47408b09e0496ec6f87d5f9ed36c1bbcb598effb Mon Sep 17 00:00:00 2001
From: Minghao Li <minghaoli@google.com>
Date: Tue, 24 Dec 2024 16:43:36 +0800
Subject: [PATCH 13/13] Polish docstring

---
 mobly/base_suite.py   | 3 ++-
 mobly/suite_runner.py | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/mobly/base_suite.py b/mobly/base_suite.py
index f2a7378a..6c4a8ef8 100644
--- a/mobly/base_suite.py
+++ b/mobly/base_suite.py
@@ -96,7 +96,8 @@ def teardown_suite(self):
     """Function used to add post tests cleanup tasks (optional)."""
     pass
 
-  # Optional interfaces to record user defined suite information to test summary
+  # Optional interfaces that users can override to record customized suite
+  # information to test summary.
 
   def get_suite_name(self):
     """Override to return a customized suite name (optional).
diff --git a/mobly/suite_runner.py b/mobly/suite_runner.py
index 67a7f13b..145fdeb3 100644
--- a/mobly/suite_runner.py
+++ b/mobly/suite_runner.py
@@ -136,6 +136,7 @@ def set_suite_name(self, suite_name):
     self._suite_name = suite_name
 
   def set_run_identifier(self, run_identifier):
+    """Sets the run identifier."""
     self._run_identifier = run_identifier
 
   def set_extras(self, extras):