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.

Comments:

Patch:
diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py
index 2be180b..5d7027e 100644
--- a/mobly/controllers/android_device.py
+++ b/mobly/controllers/android_device.py
@@ -725,7 +725,17 @@ class AndroidDevice(object):
                     ' "%s".' % (package, client_name))
         client = snippet_client.SnippetClient(
             package=package, adb_proxy=self.adb, log=self.log)
-        client.start_app_and_connect()
+        try:
+            client.start_app_and_connect()
+        except Exception as e:
+            # If errors happen, make sure we clean up before raising.
+            try:
+                client.stop_app()
+            except:
+                self.log.exception(
+                    'Failed to stop app after failure to launch.')
+            # Raise the error from start app failure.
+            raise e
         self._snippet_clients[name] = client
         setattr(self, name, client)
 
diff --git a/mobly/signals.py b/mobly/signals.py
index 85bdc30..e367a8f 100644
--- 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
index 0c540a5..24290de 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -318,7 +318,7 @@ class BaseTestClass(object):
             # 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 @@ class BaseTestClass(object):
             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 @@ class BaseTestClass(object):
         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 @@ class BaseTestClass(object):
                 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 @@ class BaseTestClass(object):
         # 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
index d615f3f..29781af 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -739,6 +739,48 @@ class BaseTestTest(unittest.TestCase):
                          ("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):