Handle SIGTERM in Mobly. (#858)

3 files changed, 72 insertions, 0 deletions

diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index b618d52..9c91477 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -16,6 +16,7 @@ import argparse
 import contextlib
 import logging
 import os
+import signal
 import sys
 import time
 
@@ -402,6 +403,19 @@ class TestRunner:
     summary_writer = records.TestSummaryWriter(
         os.path.join(self._test_run_metadata.root_output_path,
                      records.OUTPUT_FILE_SUMMARY))
+
+    # When a SIGTERM is received during the execution of a test, the Mobly test
+    # immediately terminates without executing any of the finally blocks. This
+    # handler converts the SIGTERM into a TestAbortAll signal so that the
+    # finally blocks will execute. We use TestAbortAll because other exceptions
+    # will be caught in the base test class and it will continue executing
+    # remaining tests.
+    def sigterm_handler(*args):
+      logging.warning('Test received a SIGTERM. Aborting all tests.')
+      raise signals.TestAbortAll('Test received a SIGTERM.')
+
+    signal.signal(signal.SIGTERM, sigterm_handler)
+
     try:
       for test_run_info in self._test_run_infos:
         # Set up the test-specific config
diff --git a/tests/lib/terminated_test.py b/tests/lib/terminated_test.py
new file mode 100644
index 0000000..b9dfdf4
--- /dev/null
+++ b/tests/lib/terminated_test.py
@@ -0,0 +1,38 @@
+# Copyright 2022 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.
+
+import logging
+import os
+import platform
+import signal
+
+from mobly import base_test
+from mobly import signals
+from mobly import test_runner
+
+
+class TerminatedTest(base_test.BaseTestClass):
+
+  def test_terminated(self):
+    # SIGTERM handler does not work on Windows. So just simulate the behaviour
+    # for the purpose of this test.
+    if platform.system() == 'Windows':
+      logging.warning('Test received a SIGTERM. Aborting all tests.')
+      raise signals.TestAbortAll('Test received a SIGTERM.')
+    else:
+      os.kill(os.getpid(), signal.SIGTERM)
+
+
+if __name__ == '__main__':
+  test_runner.main()
diff --git a/tests/mobly/test_runner_test.py b/tests/mobly/test_runner_test.py
index 0339c35..efddc4c 100755
--- a/tests/mobly/test_runner_test.py
+++ b/tests/mobly/test_runner_test.py
@@ -17,6 +17,7 @@ import logging
 import os
 import re
 import shutil
+import sys
 import tempfile
 import unittest
 from unittest import mock
@@ -25,12 +26,14 @@ from mobly import config_parser
 from mobly import records
 from mobly import signals
 from mobly import test_runner
+from mobly import utils
 from tests.lib import mock_android_device
 from tests.lib import mock_controller
 from tests.lib import integration_test
 from tests.lib import integration2_test
 from tests.lib import integration3_test
 from tests.lib import multiple_subclasses_module
+from tests.lib import terminated_test
 import yaml
 
 
@@ -265,6 +268,23 @@ class TestRunnerTest(unittest.TestCase):
     self.assertEqual(results['Passed'], 0)
     self.assertEqual(results['Failed'], 0)
 
+  def test_run_when_terminated(self):
+    mock_test_config = self.base_mock_test_config.copy()
+    tr = test_runner.TestRunner(self.log_dir, self.testbed_name)
+    tr.add_test_class(mock_test_config, terminated_test.TerminatedTest)
+
+    with self.assertRaises(signals.TestAbortAll):
+      with self.assertLogs(level=logging.WARNING) as log_output:
+        # Set handler log level due to bug in assertLogs.
+        # https://github.com/python/cpython/issues/86109
+        logging.getLogger().handlers[0].setLevel(logging.WARNING)
+        tr.run()
+
+    self.assertIn('Test received a SIGTERM. Aborting all tests.',
+                  log_output.output[0])
+    self.assertIn('Abort all subsequent test classes', log_output.output[1])
+    self.assertIn('Test received a SIGTERM.', log_output.output[1])
+
   def test_add_test_class_mismatched_log_path(self):
     tr = test_runner.TestRunner('/different/log/dir', self.testbed_name)
     with self.assertRaisesRegex(