Repository: google/mobly

Issue:
`signals.TestAbortClass` does not work in `setup_test` and `on_xxxx` methods
It is valid to abort the class at any stage of a test.

E.g. If we find out in `on_fail` that a test failed due to a fatal error like snippet_client losing socket connection, it is reasonable for the test to abort all subsequent tests in the class instead of continue running and failing more tests when we know they will fail.


Discussion:
The current code also has a bug: raising `TestAbortAll` in `on_fail` will cause the current record to be missed from test result.

Resolving patch:
diff --git a/mobly/signals.py b/mobly/signals.py
--- a/mobly/signals.py
+++ b/mobly/signals.py
@@ -62,13 +62,18 @@ class TestSkip(TestSignal):
     """Raised when a test has been skipped."""
 
 
-class TestAbortClass(TestSignal):
+class TestAbortSignal(TestSignal):
+    """Base class for abort signals.
+    """
+
+
+class TestAbortClass(TestAbortSignal):
     """Raised when all subsequent tests within the same test class should
     be aborted.
     """
 
 
-class TestAbortAll(TestSignal):
+class TestAbortAll(TestAbortSignal):
     """Raised when all subsequent tests should be aborted."""
 
 


Test patch:
diff --git a/mobly/base_test.py b/mobly/base_test.py
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -318,7 +318,7 @@ def _exec_procedure_func(self, func, tr_record):
             # Pass a copy of the record instead of the actual object so that it
             # will not be modified.
             func(copy.deepcopy(tr_record))
-        except signals.TestAbortAll:
+        except signals.TestAbortSignal:
             raise
         except Exception as e:
             logging.exception('Exception happened when executing %s for %s.',
@@ -363,7 +363,7 @@ def exec_one_test(self, test_name, test_method, args=(), **kwargs):
             finally:
                 try:
                     self._teardown_test(test_name)
-                except signals.TestAbortAll:
+                except signals.TestAbortSignal:
                     raise
                 except Exception as e:
                     logging.exception(e)
@@ -374,7 +374,7 @@ def exec_one_test(self, test_name, test_method, args=(), **kwargs):
         except signals.TestSkip as e:
             # Test skipped.
             tr_record.test_skip(e)
-        except (signals.TestAbortClass, signals.TestAbortAll) as e:
+        except signals.TestAbortSignal as e:
             # Abort signals, pass along.
             tr_record.test_fail(e)
             raise e
@@ -389,17 +389,20 @@ def exec_one_test(self, test_name, test_method, args=(), **kwargs):
                 tr_record.test_pass()
         finally:
             tr_record.update_record()
-            if tr_record.result in (records.TestResultEnums.TEST_RESULT_ERROR,
-                                    records.TestResultEnums.TEST_RESULT_FAIL):
-                self._exec_procedure_func(self._on_fail, tr_record)
-            elif tr_record.result == records.TestResultEnums.TEST_RESULT_PASS:
-                self._exec_procedure_func(self._on_pass, tr_record)
-            elif tr_record.result == records.TestResultEnums.TEST_RESULT_SKIP:
-                self._exec_procedure_func(self._on_skip, tr_record)
-            self.results.add_record(tr_record)
-            self.summary_writer.dump(tr_record.to_dict(),
-                                     records.TestSummaryEntryType.RECORD)
-            self.current_test_name = None
+            try:
+                if tr_record.result in (
+                        records.TestResultEnums.TEST_RESULT_ERROR,
+                        records.TestResultEnums.TEST_RESULT_FAIL):
+                    self._exec_procedure_func(self._on_fail, tr_record)
+                elif tr_record.result == records.TestResultEnums.TEST_RESULT_PASS:
+                    self._exec_procedure_func(self._on_pass, tr_record)
+                elif tr_record.result == records.TestResultEnums.TEST_RESULT_SKIP:
+                    self._exec_procedure_func(self._on_skip, tr_record)
+            finally:
+                self.results.add_record(tr_record)
+                self.summary_writer.dump(tr_record.to_dict(),
+                                         records.TestSummaryEntryType.RECORD)
+                self.current_test_name = None
 
     def _assert_function_name_in_stack(self, expected_func_name):
         """Asserts that the current stack contains the given function name."""
@@ -577,7 +580,7 @@ def run(self, test_names=None):
         # Setup for the class.
         try:
             self._setup_class()
-        except signals.TestAbortClass as e:
+        except signals.TestAbortSignal as e:
             # The test class is intentionally aborted.
             # Skip all tests peacefully.
             e.details = 'setup_class aborted due to: %s' % e.details
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -739,6 +739,48 @@ def test_3(self):
                          ("Error 0, Executed 0, Failed 0, Passed 0, "
                           "Requested 3, Skipped 3"))
 
+    def test_abort_class_in_setup_test(self):
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_test(self):
+                asserts.abort_class(MSG_EXPECTED_EXCEPTION)
+
+            def test_1(self):
+                never_call()
+
+            def test_2(self):
+                never_call()
+
+            def test_3(self):
+                never_call()
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run(test_names=["test_1", "test_2", "test_3"])
+        self.assertEqual(len(bt_cls.results.skipped), 2)
+        self.assertEqual(bt_cls.results.summary_str(),
+                         ("Error 0, Executed 1, Failed 1, Passed 0, "
+                          "Requested 3, Skipped 2"))
+
+    def test_abort_class_in_on_fail(self):
+        class MockBaseTest(base_test.BaseTestClass):
+            def test_1(self):
+                asserts.fail(MSG_EXPECTED_EXCEPTION)
+
+            def test_2(self):
+                never_call()
+
+            def test_3(self):
+                never_call()
+
+            def on_fail(self, record):
+                asserts.abort_class(MSG_EXPECTED_EXCEPTION)
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run(test_names=["test_1", "test_2", "test_3"])
+        self.assertEqual(len(bt_cls.results.skipped), 2)
+        self.assertEqual(bt_cls.results.summary_str(),
+                         ("Error 0, Executed 1, Failed 1, Passed 0, "
+                          "Requested 3, Skipped 2"))
+
     def test_setup_and_teardown_execution_count(self):
         class MockBaseTest(base_test.BaseTestClass):
             def test_func(self):