Repository: google/mobly

Issue:
`expects` APIs crash in execution stages before the test case
In `setup_class` and `setup_generated_tests`, calling `expects` APIs would cause crashes.
Also in `teardown_class`, `expects` APIs would probably attach errors to the wrong record.

Discussion:

Patch:
diff --git a/mobly/controllers/android_device_lib/adb.py b/mobly/controllers/android_device_lib/adb.py
index 41af20c..15e2a5c 100644
--- a/mobly/controllers/android_device_lib/adb.py
+++ b/mobly/controllers/android_device_lib/adb.py
@@ -44,13 +44,17 @@ class AdbError(Error):
         stdout: byte string, the raw stdout of the command.
         stderr: byte string, the raw stderr of the command.
         ret_code: int, the return code of the command.
+        serial: string, the serial of the device the command is executed on.
+            This is an empty string if the adb command is not specific to a
+            device.
     """
 
-    def __init__(self, cmd, stdout, stderr, ret_code):
+    def __init__(self, cmd, stdout, stderr, ret_code, serial=''):
         self.cmd = cmd
         self.stdout = stdout
         self.stderr = stderr
         self.ret_code = ret_code
+        self.serial = serial
 
     def __str__(self):
         return ('Error executing adb cmd "%s". ret: %d, stdout: %s, stderr: %s'
@@ -64,11 +68,15 @@ class AdbTimeoutError(Error):
     Args:
         cmd: list of strings, the adb command that timed out
         timeout: float, the number of seconds passed before timing out.
+        serial: string, the serial of the device the command is executed on.
+            This is an empty string if the adb command is not specific to a
+            device.
     """
 
-    def __init__(self, cmd, timeout):
+    def __init__(self, cmd, timeout, serial=''):
         self.cmd = cmd
         self.timeout = timeout
+        self.serial = serial
 
     def __str__(self):
         return 'Timed out executing command "%s" after %ss.' % (
@@ -155,6 +163,7 @@ class AdbProxy(object):
             The output of the adb command run if exit code is 0.
 
         Raises:
+            ValueError: timeout value is invalid.
             AdbError: The adb command exit code is not 0.
             AdbTimeoutError: The adb command timed out.
         """
@@ -162,13 +171,14 @@ class AdbProxy(object):
             args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell)
         process = psutil.Process(proc.pid)
         if timeout and timeout <= 0:
-            raise Error('Timeout is not a positive value: %s' % timeout)
+            raise ValueError('Timeout is not a positive value: %s' % timeout)
         if timeout and timeout > 0:
             try:
                 process.wait(timeout=timeout)
             except psutil.TimeoutExpired:
                 process.terminate()
-                raise AdbTimeoutError(cmd=args, timeout=timeout)
+                raise AdbTimeoutError(
+                    cmd=args, timeout=timeout, serial=self.serial)
 
         (out, err) = proc.communicate()
         if stderr:
@@ -179,7 +189,12 @@ class AdbProxy(object):
         if ret == 0:
             return out
         else:
-            raise AdbError(cmd=args, stdout=out, stderr=err, ret_code=ret)
+            raise AdbError(
+                cmd=args,
+                stdout=out,
+                stderr=err,
+                ret_code=ret,
+                serial=self.serial)
 
     def _execute_and_process_stdout(self, args, shell, handler):
         """Executes adb commands and processes the stdout with a handler.
diff --git a/mobly/utils.py b/mobly/utils.py
index a9f065c..bb2f316 100644
--- a/mobly/utils.py
+++ b/mobly/utils.py
@@ -283,7 +283,7 @@ def concurrent_exec(func, param_list):
         return return_vals
 
 
-def start_standing_subprocess(cmd, shell=False):
+def start_standing_subprocess(cmd, shell=False, env=None):
     """Starts a long-running subprocess.
 
     This is not a blocking call and the subprocess started by it should be
@@ -296,6 +296,9 @@ def start_standing_subprocess(cmd, shell=False):
         cmd: string, the command to start the subprocess with.
         shell: bool, True to run this command through the system shell,
             False to invoke it directly. See subprocess.Proc() docs.
+        env: dict, a custom environment to run the standing subprocess. If not
+            specified, inherits the current environment. See subprocess.Popen()
+            docs.
 
     Returns:
         The subprocess that was started.
@@ -306,7 +309,8 @@ def start_standing_subprocess(cmd, shell=False):
         stdin=subprocess.PIPE,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
-        shell=shell)
+        shell=shell,
+        env=env)
     # Leaving stdin open causes problems for input, e.g. breaking the
     # code.inspect() shell (http://stackoverflow.com/a/25512460/1612937), so
     # explicitly close it assuming it is not needed for standing subprocesses.


Test patch:
diff --git a/mobly/base_test.py b/mobly/base_test.py
index 03db500..56d8d4e 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -292,9 +292,46 @@ class BaseTestClass(object):
     def _setup_class(self):
         """Proxy function to guarantee the base implementation of setup_class
         is called.
+
+        Returns:
+            If `self.results` is returned instead of None, this means something
+            has gone wrong, and the rest of the test class should not execute.
         """
-        with self._log_test_stage(STAGE_NAME_SETUP_CLASS):
-            self.setup_class()
+        # Setup for the class.
+        class_record = records.TestResultRecord(STAGE_NAME_SETUP_CLASS,
+                                                self.TAG)
+        class_record.test_begin()
+        self.current_test_info = runtime_test_info.RuntimeTestInfo(
+            STAGE_NAME_SETUP_CLASS, self.log_path, class_record)
+        expects.recorder.reset_internal_states(class_record)
+        try:
+            with self._log_test_stage(STAGE_NAME_SETUP_CLASS):
+                self.setup_class()
+        except signals.TestAbortSignal:
+            # Throw abort signals to outer try block for handling.
+            raise
+        except Exception as e:
+            # Setup class failed for unknown reasons.
+            # Fail the class and skip all tests.
+            logging.exception('Error in %s#setup_class.', self.TAG)
+            class_record.test_error(e)
+            self.results.add_class_error(class_record)
+            self._exec_procedure_func(self._on_fail, class_record)
+            class_record.update_record()
+            self.summary_writer.dump(class_record.to_dict(),
+                                     records.TestSummaryEntryType.RECORD)
+            self._skip_remaining_tests(e)
+            return self.results
+        if expects.recorder.has_error:
+            self._exec_procedure_func(self._on_fail, class_record)
+            class_record.test_error()
+            class_record.update_record()
+            self.summary_writer.dump(class_record.to_dict(),
+                                     records.TestSummaryEntryType.RECORD)
+            self.results.add_class_error(class_record)
+            self._skip_remaining_tests(
+                class_record.termination_signal.exception)
+            return self.results
 
     def setup_class(self):
         """Setup function that will be called before executing any test in the
@@ -302,6 +339,8 @@ class BaseTestClass(object):
 
         To signal setup failure, use asserts or raise your own exception.
 
+        Errors raised from `setup_class` will trigger `on_fail`.
+
         Implementation is optional.
         """
 
@@ -314,6 +353,7 @@ class BaseTestClass(object):
         record.test_begin()
         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.teardown_class()
@@ -327,6 +367,12 @@ class BaseTestClass(object):
             self.results.add_class_error(record)
             self.summary_writer.dump(record.to_dict(),
                                      records.TestSummaryEntryType.RECORD)
+        else:
+            if expects.recorder.has_error:
+                record.update_record()
+                self.results.add_class_error(record)
+                self.summary_writer.dump(record.to_dict(),
+                                         records.TestSummaryEntryType.RECORD)
         finally:
             # Write controller info and summary to summary file.
             self._record_controller_info()
@@ -336,6 +382,8 @@ class BaseTestClass(object):
         """Teardown function that will be called after all the selected tests in
         the test class have been executed.
 
+        Errors raised from `teardown_class` do not trigger `on_fail`.
+
         Implementation is optional.
         """
 
@@ -775,28 +823,9 @@ class BaseTestClass(object):
                                  records.TestSummaryEntryType.TEST_NAME_LIST)
         tests = self._get_test_methods(test_names)
         try:
-            # Setup for the class.
-            class_record = records.TestResultRecord(STAGE_NAME_SETUP_CLASS,
-                                                    self.TAG)
-            class_record.test_begin()
-            self.current_test_info = runtime_test_info.RuntimeTestInfo(
-                STAGE_NAME_SETUP_CLASS, self.log_path, class_record)
-            try:
-                self._setup_class()
-            except signals.TestAbortSignal:
-                # Throw abort signals to outer try block for handling.
-                raise
-            except Exception as e:
-                # Setup class failed for unknown reasons.
-                # Fail the class and skip all tests.
-                logging.exception('Error in %s#setup_class.', self.TAG)
-                class_record.test_error(e)
-                self.results.add_class_error(class_record)
-                self.summary_writer.dump(class_record.to_dict(),
-                                         records.TestSummaryEntryType.RECORD)
-                self._exec_procedure_func(self._on_fail, class_record)
-                self._skip_remaining_tests(e)
-                return self.results
+            setup_class_result = self._setup_class()
+            if setup_class_result:
+                return setup_class_result
             # Run tests in order.
             for test_name, test_method in tests:
                 self.exec_one_test(test_name, test_method)
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index 7fb2b1d..1183133 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -1485,6 +1485,122 @@ class BaseTestTest(unittest.TestCase):
         must_call.assert_called_with('ha')
         self.assertEqual(len(bt_cls.results.passed), 2)
 
+    @mock.patch('mobly.records.TestSummaryWriter.dump')
+    def test_expect_in_setup_class(self, mock_dump):
+        must_call = mock.Mock()
+        must_call2 = mock.Mock()
+
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_class(self):
+                expects.expect_true(
+                    False, MSG_EXPECTED_EXCEPTION, extras=MOCK_EXTRA)
+                must_call('ha')
+
+            def test_func(self):
+                pass
+
+            def on_fail(self, record):
+                must_call2('on_fail')
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        must_call.assert_called_once_with('ha')
+        must_call2.assert_called_once_with('on_fail')
+        actual_record = bt_cls.results.error[0]
+        self.assertEqual(actual_record.test_name, 'setup_class')
+        self.assertEqual(actual_record.details, MSG_EXPECTED_EXCEPTION)
+        self.assertEqual(actual_record.extras, MOCK_EXTRA)
+        # Verify the class record is written out correctly.
+        setup_class_dict = mock_dump.call_args_list[1][0][0]
+        self.assertIsNotNone(setup_class_dict['Begin Time'])
+        self.assertIsNotNone(setup_class_dict['End Time'])
+        self.assertEqual(setup_class_dict['Test Name'], 'setup_class')
+
+    @mock.patch('mobly.records.TestSummaryWriter.dump')
+    def test_expect_in_setup_class_and_on_fail(self, mock_dump):
+        must_call = mock.Mock()
+        must_call2 = mock.Mock()
+
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_class(self):
+                expects.expect_true(
+                    False, 'Failure in setup_class', extras=MOCK_EXTRA)
+                must_call('ha')
+
+            def test_func(self):
+                pass
+
+            def on_fail(self, record):
+                expects.expect_true(
+                    False, 'Failure in on_fail', extras=MOCK_EXTRA)
+                must_call2('on_fail')
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        must_call.assert_called_once_with('ha')
+        must_call2.assert_called_once_with('on_fail')
+        actual_record = bt_cls.results.error[0]
+        self.assertEqual(actual_record.test_name, 'setup_class')
+        self.assertEqual(actual_record.details, 'Failure in setup_class')
+        self.assertEqual(actual_record.extras, MOCK_EXTRA)
+        on_fail_error = next(iter(actual_record.extra_errors.values()))
+        self.assertEqual(on_fail_error.details, 'Failure in on_fail')
+        self.assertEqual(on_fail_error.extras, MOCK_EXTRA)
+        # Verify the class record is written out correctly.
+        setup_class_dict = mock_dump.call_args_list[1][0][0]
+        self.assertIsNotNone(setup_class_dict['Begin Time'])
+        self.assertIsNotNone(setup_class_dict['End Time'])
+        self.assertEqual(setup_class_dict['Test Name'], 'setup_class')
+        # Verify the on_fail error is recorded in summary result.
+        extra_error_dict = next(
+            iter(setup_class_dict['Extra Errors'].values()))
+        self.assertEqual(extra_error_dict['Details'], 'Failure in on_fail')
+
+    def test_expect_in_teardown_class(self):
+        must_call = mock.Mock()
+
+        class MockBaseTest(base_test.BaseTestClass):
+            def test_func(self):
+                pass
+
+            def teardown_class(self):
+                expects.expect_true(
+                    False, MSG_EXPECTED_EXCEPTION, extras=MOCK_EXTRA)
+                must_call('ha')
+
+        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, 'teardown_class')
+        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()
+
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_test(self):
+                expects.expect_true(
+                    False, MSG_EXPECTED_EXCEPTION, extras=MOCK_EXTRA)
+                must_call('ha')
+
+            def test_func(self):
+                pass
+
+            def on_fail(self, record):
+                must_call2('on_fail')
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        must_call.assert_called_once_with('ha')
+        must_call2.assert_called_once_with('on_fail')
+        actual_record = bt_cls.results.failed[0]
+        self.assertEqual(actual_record.test_name, 'test_func')
+        self.assertEqual(actual_record.details, MSG_EXPECTED_EXCEPTION)
+        self.assertEqual(actual_record.extras, MOCK_EXTRA)
+
     def test_expect_in_teardown_test(self):
         must_call = mock.Mock()
         must_call2 = mock.Mock()
diff --git a/tests/mobly/controllers/android_device_lib/adb_test.py b/tests/mobly/controllers/android_device_lib/adb_test.py
index 71fab9c..9430795 100755
--- a/tests/mobly/controllers/android_device_lib/adb_test.py
+++ b/tests/mobly/controllers/android_device_lib/adb_test.py
@@ -91,15 +91,29 @@ class AdbTest(unittest.TestCase):
 
     @mock.patch('mobly.controllers.android_device_lib.adb.subprocess.Popen')
     @mock.patch('mobly.controllers.android_device_lib.adb.psutil.Process')
-    def test_exec_cmd_error_no_timeout(self, mock_psutil_process, mock_popen):
+    def test_exec_cmd_error_with_serial(self, mock_psutil_process, mock_popen):
         self._mock_process(mock_psutil_process, mock_popen)
         # update return code to indicate command execution error
         mock_popen.return_value.returncode = 1
+        mock_serial = 'ABCD1234'
+        with self.assertRaisesRegex(adb.AdbError,
+                                    'Error executing adb cmd .*') as context:
+            adb.AdbProxy(mock_serial).fake_cmd()
+        self.assertEqual(context.exception.serial, mock_serial)
+        self.assertIn(mock_serial, context.exception.cmd)
 
+    @mock.patch('mobly.controllers.android_device_lib.adb.subprocess.Popen')
+    @mock.patch('mobly.controllers.android_device_lib.adb.psutil.Process')
+    def test_exec_cmd_error_without_serial(self, mock_psutil_process,
+                                           mock_popen):
+        self._mock_process(mock_psutil_process, mock_popen)
+        # update return code to indicate command execution error
+        mock_popen.return_value.returncode = 1
         with self.assertRaisesRegex(adb.AdbError,
-                                    'Error executing adb cmd .*'):
+                                    'Error executing adb cmd .*') as context:
             adb.AdbProxy()._exec_cmd(
                 ['fake_cmd'], shell=False, timeout=None, stderr=None)
+        self.assertFalse(context.exception.serial)
 
     @mock.patch('mobly.controllers.android_device_lib.adb.subprocess.Popen')
     @mock.patch('mobly.controllers.android_device_lib.adb.psutil.Process')
@@ -115,23 +129,34 @@ class AdbTest(unittest.TestCase):
     @mock.patch('mobly.controllers.android_device_lib.adb.psutil.Process')
     def test_exec_cmd_timed_out(self, mock_psutil_process, mock_popen):
         self._mock_process(mock_psutil_process, mock_popen)
-        # mock process.wait(timeout=timeout) to
-        # throw psutil.TimeoutExpired exception
         mock_psutil_process.return_value.wait.side_effect = (
             adb.psutil.TimeoutExpired('Timed out'))
+        mock_serial = '1234Abcd'
+        with self.assertRaisesRegex(
+                adb.AdbTimeoutError, 'Timed out executing command "adb -s '
+                '1234Abcd fake-cmd" after 0.01s.') as context:
+            adb.AdbProxy(mock_serial).fake_cmd(timeout=0.01)
+        self.assertEqual(context.exception.serial, mock_serial)
+        self.assertIn(mock_serial, context.exception.cmd)
 
+    @mock.patch('mobly.controllers.android_device_lib.adb.subprocess.Popen')
+    @mock.patch('mobly.controllers.android_device_lib.adb.psutil.Process')
+    def test_exec_cmd_timed_out_without_serial(self, mock_psutil_process,
+                                               mock_popen):
+        self._mock_process(mock_psutil_process, mock_popen)
+        mock_psutil_process.return_value.wait.side_effect = (
+            adb.psutil.TimeoutExpired('Timed out'))
         with self.assertRaisesRegex(adb.AdbTimeoutError,
-                                    'Timed out executing command "fake_cmd" '
-                                    'after 0.1s.'):
-            adb.AdbProxy()._exec_cmd(
-                ['fake_cmd'], shell=False, timeout=0.1, stderr=None)
+                                    'Timed out executing command "adb '
+                                    'fake-cmd" after 0.01s.') as context:
+            adb.AdbProxy().fake_cmd(timeout=0.01)
 
     @mock.patch('mobly.controllers.android_device_lib.adb.subprocess.Popen')
     @mock.patch('mobly.controllers.android_device_lib.adb.psutil.Process')
     def test_exec_cmd_with_negative_timeout_value(self, mock_psutil_process,
                                                   mock_popen):
         self._mock_process(mock_psutil_process, mock_popen)
-        with self.assertRaisesRegex(adb.Error,
+        with self.assertRaisesRegex(ValueError,
                                     'Timeout is not a positive value: -1'):
             adb.AdbProxy()._exec_cmd(
                 ['fake_cmd'], shell=False, timeout=-1, stderr=None)
@@ -156,8 +181,8 @@ class AdbTest(unittest.TestCase):
         self._mock_execute_and_process_stdout_process(mock_popen)
         mock_handler = mock.MagicMock()
         mock_popen.return_value.communicate = mock.Mock(
-            return_value=(unexpected_stdout, MOCK_DEFAULT_STDERR.encode(
-                'utf-8')))
+            return_value=(unexpected_stdout,
+                          MOCK_DEFAULT_STDERR.encode('utf-8')))
 
         err = adb.AdbProxy()._execute_and_process_stdout(
             ['fake_cmd'], shell=False, handler=mock_handler)
diff --git a/tests/mobly/utils_test.py b/tests/mobly/utils_test.py
index 605dc8c..b9b0c22 100755
--- a/tests/mobly/utils_test.py
+++ b/tests/mobly/utils_test.py
@@ -53,6 +53,31 @@ class UtilsTest(unittest.TestCase):
             p.stderr.close()
             p.wait()
 
+    @mock.patch('subprocess.Popen')
+    def test_start_standing_subproc_without_env(self, mock_Popen):
+        p = utils.start_standing_subprocess(self.sleep_cmd)
+        mock_Popen.assert_called_with(
+            self.sleep_cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            shell=False,
+            env=None,
+        )
+
+    @mock.patch('subprocess.Popen')
+    def test_start_standing_subproc_with_custom_env(self, mock_Popen):
+        mock_env = mock.MagicMock(spec=dict)
+        p = utils.start_standing_subprocess(self.sleep_cmd, env=mock_env)
+        mock_Popen.assert_called_with(
+            self.sleep_cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            shell=False,
+            env=mock_env,
+        )
+
     def test_stop_standing_subproc(self):
         p = utils.start_standing_subprocess([self.sleep_cmd, '4'])
         p1 = psutil.Process(p.pid)