Record the class termination signal in TestClassInfoRecord

mobly/base_test.py

@@ -376,9 +376,13 @@ def _pre_run(self):
     self.current_test_info = runtime_test_info.RuntimeTestInfo(
         stage_name, self.log_path, record
     )
+    expects.recorder.reset_internal_states(record)
     try:
       with self._log_test_stage(stage_name):
         self.pre_run()
+      if expects.recorder.has_error:
+        self._record_class_error(record)
+        return False
       return True
     except _UNCATCHABLE_EXCEPTIONS:
       raise
@@ -1180,6 +1184,7 @@ def run(self, test_names=None):
         records.TestSummaryEntryType.TEST_NAME_LIST,
     )
     tests = self._get_test_methods(test_names)
+    class_termination_error = None
     try:
       setup_class_result = self._setup_class()
       if setup_class_result:
@@ -1239,24 +1244,37 @@ def run(self, test_names=None):
       return self.results
     except signals.TestAbortClass as e:
       e.details = f'Test class aborted due to: {e.details}'
+      class_termination_error = e
       self._skip_remaining_tests(e)
       return self.results
     except signals.TestAbortAll as e:
       e.details = f'All remaining tests aborted due to: {e.details}'
+      class_termination_error = e
       self._skip_remaining_tests(e)
       # Piggy-back test results on this exception object so we don't lose
       # results from this test class.
       setattr(e, 'results', self.results)
       raise e
+    except BaseException as e:
+      class_termination_error = e
+      raise
     finally:
-      self._teardown_class()
-      class_info.class_end()
-      self.summary_writer.dump(
-          class_info.to_dict(), records.TestSummaryEntryType.TEST_CLASS_INFO
-      )
-      logging.info(
-          'Summary for test class %s: %s', self.TAG, self.results.summary_str()
-      )
+      try:
+        self._teardown_class()
+      except BaseException as e:
+        # Whatever escapes teardown is what the caller sees terminate the class.
+        class_termination_error = e
+        raise
+      finally:
+        class_info.class_end(class_termination_error)
+        self.summary_writer.dump(
+            class_info.to_dict(), records.TestSummaryEntryType.TEST_CLASS_INFO
+        )
+        logging.info(
+            'Summary for test class %s: %s',
+            self.TAG,
+            self.results.summary_str(),
+        )
 
   def _clean_up(self):
     """The final stage of a test class execution."""

mobly/records.py

@@ -227,32 +227,47 @@ class TestClassInfoRecord:
 
   Captures the wall-clock span of the whole class, including `setup_class`
   and `teardown_class`, which is not derivable from the individual test
-  records. Times are epoch milliseconds, consistent with `TestResultRecord`.
+  records, and the signal that terminated the class early, if any. Times are
+  epoch milliseconds, consistent with `TestResultRecord`.
   """
 
-  KEY_TEST_CLASS = TestResultEnums.RECORD_CLASS
-  KEY_BEGIN_TIME = TestResultEnums.RECORD_BEGIN_TIME
-  KEY_END_TIME = TestResultEnums.RECORD_END_TIME
-
   def __init__(self, test_class):
     self.test_class = test_class
     self.begin_time = None
     self.end_time = None
+    self.termination_signal = None
 
   def class_begin(self):
     """Marks the beginning of the test class execution."""
     self.begin_time = utils.get_current_epoch_time()
 
-  def class_end(self):
-    """Marks the end of the test class execution."""
+  def class_end(self, e=None):
+    """Marks the end of the test class execution.
+
+    Args:
+      e: The exception or `TestSignal` that terminated the class early, if any.
+    """
     self.end_time = utils.get_current_epoch_time()
+    if e is not None:
+      self.termination_signal = ExceptionRecord(e)
 
   def to_dict(self):
-    result = {}
-    result[self.KEY_TEST_CLASS] = self.test_class
-    result[self.KEY_BEGIN_TIME] = self.begin_time
-    result[self.KEY_END_TIME] = self.end_time
-    return result
+    signal = self.termination_signal
+    return {
+        TestResultEnums.RECORD_CLASS: self.test_class,
+        TestResultEnums.RECORD_BEGIN_TIME: self.begin_time,
+        TestResultEnums.RECORD_END_TIME: self.end_time,
+        TestResultEnums.RECORD_DETAILS: signal.details if signal else None,
+        TestResultEnums.RECORD_TERMINATION_SIGNAL_TYPE: (
+            signal.type if signal else None
+        ),
+        TestResultEnums.RECORD_STACKTRACE: (
+            signal.stacktrace if signal else None
+        ),
+        TestResultEnums.RECORD_EXTRAS: (
+            copy.deepcopy(signal.extras) if signal else None
+        ),
+    }
 
   def __repr__(self):
     return str(self.to_dict())

tests/mobly/base_test_test.py

@@ -1126,7 +1126,8 @@ def test_something(self):
     )
     self.assertEqual(bt_cls.results.summary_str(), expected_summary)
 
-  def test_abort_class_setup_class(self):
+  @mock.patch('mobly.records.TestSummaryWriter.dump')
+  def test_abort_class_setup_class(self, mock_dump):
     """A class was intentionally aborted by the test.
 
     This is not considered an error as the abort class is used as a skip
@@ -1155,6 +1156,16 @@ def test_3(self):
         bt_cls.results.summary_str(),
         'Error 0, Executed 0, Failed 0, Passed 0, Requested 3, Skipped 3',
     )
+    # The abort signal is recorded as the class's termination signal.
+    class_info = next(
+        call.args[0]
+        for call in mock_dump.call_args_list
+        if call.args[1] == records.TestSummaryEntryType.TEST_CLASS_INFO
+    )
+    self.assertEqual(
+        class_info[records.TestResultEnums.RECORD_TERMINATION_SIGNAL_TYPE],
+        'TestAbortClass',
+    )
 
   def test_abort_class_in_setup_test(self):
     class MockBaseTest(base_test.BaseTestClass):
@@ -2100,6 +2111,26 @@ def teardown_class(self):
     self.assertEqual(actual_record.extras, MOCK_EXTRA)
     self.assertIsNotNone(actual_record.end_time)
 
+  def test_expect_in_pre_run(self):
+    must_call = mock.Mock()
+
+    class MockBaseTest(base_test.BaseTestClass):
+
+      def pre_run(self):
+        expects.expect_true(False, MSG_EXPECTED_EXCEPTION, extras=MOCK_EXTRA)
+        must_call('ha')
+
+      def setup_class(self):
+        never_call()
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    bt_cls.run()
+    must_call.assert_called_once_with('ha')
+    actual_record = bt_cls.results.error[0]
+    self.assertEqual(actual_record.test_name, 'pre_run')
+    self.assertEqual(actual_record.details, MSG_EXPECTED_EXCEPTION)
+    self.assertEqual(actual_record.extras, MOCK_EXTRA)
+
   def test_expect_in_setup_test(self):
     must_call = mock.Mock()
     must_call2 = mock.Mock()

tests/mobly/records_test.py

@@ -567,6 +567,10 @@ def test_test_class_info_record(self):
             records.TestResultEnums.RECORD_CLASS: 'SomeTestClass',
             records.TestResultEnums.RECORD_BEGIN_TIME: None,
             records.TestResultEnums.RECORD_END_TIME: None,
+            records.TestResultEnums.RECORD_DETAILS: None,
+            records.TestResultEnums.RECORD_TERMINATION_SIGNAL_TYPE: None,
+            records.TestResultEnums.RECORD_STACKTRACE: None,
+            records.TestResultEnums.RECORD_EXTRAS: None,
         },
     )
     record.class_begin()
@@ -581,6 +585,22 @@ def test_test_class_info_record(self):
     # Must be serializable by the summary writer.
     yaml.safe_dump(d)
 
+  def test_test_class_info_record_with_termination_signal(self):
+    record = records.TestClassInfoRecord('SomeTestClass')
+    record.class_begin()
+    record.class_end(
+        signals.TestAbortClass('device not supported', extras={'k': 'v'})
+    )
+    d = record.to_dict()
+    self.assertEqual(
+        d[records.TestResultEnums.RECORD_DETAILS], 'device not supported'
+    )
+    self.assertEqual(
+        d[records.TestResultEnums.RECORD_TERMINATION_SIGNAL_TYPE],
+        'TestAbortClass',
+    )
+    self.assertEqual(d[records.TestResultEnums.RECORD_EXTRAS], {'k': 'v'})
+
 
 if __name__ == '__main__':
   unittest.main()