Snap for 10909147 from a529b98b550f5d2cdb797699d6528ca5454759a6 to 24Q1-releaseandroid-14.0.0_r37android-14.0.0_r36android-14.0.0_r35android-14.0.0_r34android-14.0.0_r33android-14.0.0_r32android-14.0.0_r31android-14.0.0_r30android-14.0.0_r29android14-qpr2-s5-releaseandroid14-qpr2-s4-releaseandroid14-qpr2-s3-releaseandroid14-qpr2-s2-releaseandroid14-qpr2-s1-releaseandroid14-qpr2-release

Change-Id: Ic987d2a361cb62dfd6bc337ca88142a8adaf5edd

16 files changed, 218 insertions, 28 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a94d23c..42c8df7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,20 @@
 # Mobly Release History
 
 
+## Mobly Release 1.12.2: Improve Support for Custom Suites
+
+Bug fixes and improvements to better support users who construct their own
+suite based on `test_runner` APIs and `suite_runner`.
+
+### Fixes
+* Make print test case name feature usable.
+* Ensure default log path exists.
+* Missing info in test records are now populated.
+* Enable Android devices in bootloader mode to be picked up in registration.
+
+[Full list of changes](https://github.com/google/mobly/milestone/29?closed=1)
+
+
 ## Mobly Release 1.12.1: Minor Improvements and Fixes
 
 ### New
diff --git a/mobly/base_test.py b/mobly/base_test.py
index e7da22c..9bf4fcc 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -170,6 +170,12 @@ class BaseTestClass:
       the test logic.
   """
 
+  # Explicitly set the type since we set this to `None` in between
+  # test cases executions when there's no active test. However, since
+  # it is safe for clients to call at any point during normal execution
+  # of a Mobly test, we avoid using the `Optional` type hint for convenience.
+  current_test_info: runtime_test_info.RuntimeTestInfo
+
   TAG = None
 
   def __init__(self, configs):
@@ -797,6 +803,7 @@ class BaseTestClass:
         else:
           # Check if anything failed by `expects`.
           if before_count < expects.recorder.error_count:
+            tr_record.test_error()
             teardown_test_failed = True
     except (signals.TestFailure, AssertionError) as e:
       tr_record.test_fail(e)
diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py
index 54166a5..c093e3c 100644
--- a/mobly/controllers/android_device.py
+++ b/mobly/controllers/android_device.py
@@ -901,7 +901,7 @@ class AndroidDevice:
     # So we need to wait for the device to come back before proceeding.
     self.adb.wait_for_device(timeout=DEFAULT_TIMEOUT_BOOT_COMPLETION_SECOND)
 
-  def load_snippet(self, name, package):
+  def load_snippet(self, name, package, config=None):
     """Starts the snippet apk with the given package name and connects.
 
     Examples:
@@ -917,6 +917,9 @@ class AndroidDevice:
         client. E.g. `name='maps'` attaches the snippet client to
         `ad.maps`.
       package: string, the package name of the snippet apk to connect to.
+      config: snippet_client_v2.Config, the configuration object for
+        controlling the snippet behaviors. See the docstring of the `Config`
+        class for supported configurations.
 
     Raises:
       SnippetError: Illegal load operations are attempted.
@@ -926,7 +929,9 @@ class AndroidDevice:
       raise SnippetError(
           self,
           'Attribute "%s" already exists, please use a different name.' % name)
-    self.services.snippets.add_snippet_client(name, package)
+    self.services.snippets.add_snippet_client(
+        name, package, config=config
+    )
 
   def unload_snippet(self, name):
     """Stops a snippet apk.
diff --git a/mobly/controllers/android_device_lib/services/logcat.py b/mobly/controllers/android_device_lib/services/logcat.py
index cbd8e95..37a8454 100644
--- a/mobly/controllers/android_device_lib/services/logcat.py
+++ b/mobly/controllers/android_device_lib/services/logcat.py
@@ -201,6 +201,10 @@ class Logcat(base_service.BaseService):
 
     The collection runs in a separate subprocess and saves logs in a file.
     """
+    if self._ad.is_bootloader:
+      self._ad.log.warning(
+          'Skip starting logcat because the device is in fastboot mode.')
+      return
     self._assert_not_running()
     if self._config.clear_log:
       self.clear_adb_log()
diff --git a/mobly/controllers/android_device_lib/services/snippet_management_service.py b/mobly/controllers/android_device_lib/services/snippet_management_service.py
index fae60e2..05e8cda 100644
--- a/mobly/controllers/android_device_lib/services/snippet_management_service.py
+++ b/mobly/controllers/android_device_lib/services/snippet_management_service.py
@@ -55,7 +55,7 @@ class SnippetManagementService(base_service.BaseService):
     if name in self._snippet_clients:
       return self._snippet_clients[name]
 
-  def add_snippet_client(self, name, package):
+  def add_snippet_client(self, name, package, config=None):
     """Adds a snippet client to the management.
 
     Args:
@@ -63,6 +63,9 @@ class SnippetManagementService(base_service.BaseService):
         client. E.g. `name='maps'` attaches the snippet client to
         `ad.maps`.
       package: string, the package name of the snippet apk to connect to.
+      config: snippet_client_v2.Config, the configuration object for
+        controlling the snippet behaviors. See the docstring of the `Config`
+        class for supported configurations.
 
     Raises:
       Error, if a duplicated name or package is passed in.
@@ -79,7 +82,9 @@ class SnippetManagementService(base_service.BaseService):
             self, 'Snippet package "%s" has already been loaded under name'
             ' "%s".' % (package, snippet_name))
 
-    client = snippet_client_v2.SnippetClientV2(package=package, ad=self._device)
+    client = snippet_client_v2.SnippetClientV2(
+        package=package, ad=self._device, config=config,
+    )
     client.initialize()
     self._snippet_clients[name] = client
 
diff --git a/mobly/controllers/android_device_lib/snippet_client_v2.py b/mobly/controllers/android_device_lib/snippet_client_v2.py
index 3adfde5..f7494c2 100644
--- a/mobly/controllers/android_device_lib/snippet_client_v2.py
+++ b/mobly/controllers/android_device_lib/snippet_client_v2.py
@@ -13,10 +13,12 @@
 # limitations under the License.
 """Snippet Client V2 for Interacting with Snippet Server on Android Device."""
 
+import dataclasses
 import enum
 import json
 import re
 import socket
+from typing import Dict
 
 from mobly import utils
 from mobly.controllers.android_device_lib import adb
@@ -30,8 +32,8 @@ _INSTRUMENTATION_RUNNER_PACKAGE = 'com.google.android.mobly.snippet.SnippetRunne
 
 # The command template to start the snippet server
 _LAUNCH_CMD = (
-    '{shell_cmd} am instrument {user} -w -e action start {snippet_package}/'
-    f'{_INSTRUMENTATION_RUNNER_PACKAGE}')
+    '{shell_cmd} am instrument {user} -w -e action start {instrument_options} '
+    f'{{snippet_package}}/{_INSTRUMENTATION_RUNNER_PACKAGE}')
 
 # The command template to stop the snippet server
 _STOP_CMD = ('am instrument {user} -w -e action stop {snippet_package}/'
@@ -76,6 +78,23 @@ _SOCKET_READ_TIMEOUT = 60 * 10
 _CALLBACK_DEFAULT_TIMEOUT_SEC = 60 * 2
 
 
+@dataclasses.dataclass
+class Config:
+  """A configuration class.
+
+  Attributes:
+    am_instrument_options: The Android am instrument options used for
+      controlling the `onCreate` process of the app under test. Note that this
+      should only be used for controlling the app launch process, options for
+      other purposes may not take effect and you should use snippet RPCs. This
+      is because Mobly snippet runner changes the subsequent instrumentation
+      process.
+  """
+
+  am_instrument_options: Dict[str, str] = dataclasses.field(
+      default_factory=dict)
+
+
 class ConnectionHandshakeCommand(enum.Enum):
   """Commands to send to the server when sending the handshake request.
 
@@ -109,12 +128,14 @@ class SnippetClientV2(client_base.ClientBase):
       the connection to the server is made successfully.
   """
 
-  def __init__(self, package, ad):
+  def __init__(self, package, ad, config=None):
     """Initializes the instance of Snippet Client V2.
 
     Args:
       package: str, see base class.
       ad: AndroidDevice, the android device object associated with this client.
+      config: Config, the configuration object. See the docstring of the
+        `Config` class for supported configurations.
     """
     super().__init__(package=package, device=ad)
     self.host_port = None
@@ -126,6 +147,7 @@ class SnippetClientV2(client_base.ClientBase):
     self._client = None  # keep it to prevent close errors on connect failure
     self._conn = None
     self._event_client = None
+    self._config = config or Config()
 
   @property
   def user_id(self):
@@ -231,9 +253,11 @@ class SnippetClientV2(client_base.ClientBase):
     self.log.debug('Snippet server for package %s is using protocol %d.%d',
                    self.package, _PROTOCOL_MAJOR_VERSION,
                    _PROTOCOL_MINOR_VERSION)
+    option_str = self._get_instrument_options_str()
     cmd = _LAUNCH_CMD.format(shell_cmd=persists_shell_cmd,
                              user=self._get_user_command_string(),
-                             snippet_package=self.package)
+                             snippet_package=self.package,
+                             instrument_options=option_str)
     self._proc = self._run_adb_cmd(cmd)
 
     # Check protocol version and get the device port
@@ -272,6 +296,19 @@ class SnippetClientV2(client_base.ClientBase):
         _SETSID_COMMAND, _NOHUP_COMMAND)
     return ''
 
+  def _get_instrument_options_str(self):
+    self.log.debug(
+        'Got am instrument options in snippet client for package %s: %s',
+        self.package,
+        self._config.am_instrument_options,
+    )
+    if not self._config.am_instrument_options:
+      return ''
+
+    return ' '.join(
+        f'-e {k} {v}' for k, v in self._config.am_instrument_options.items()
+    )
+
   def _get_user_command_string(self):
     """Gets the appropriate command argument for specifying device user ID.
 
diff --git a/mobly/controllers/attenuator.py b/mobly/controllers/attenuator.py
index cbe3cb3..eebc0e4 100644
--- a/mobly/controllers/attenuator.py
+++ b/mobly/controllers/attenuator.py
@@ -129,6 +129,7 @@ class AttenuatorPath:
     self.model = attenuation_device.model
     self.attenuation_device = attenuation_device
     self.idx = idx
+    self.name = name
     if (self.idx >= attenuation_device.path_count):
       raise IndexError("Attenuator index out of range!")
 
diff --git a/mobly/records.py b/mobly/records.py
index b77817c..08ebe8e 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -321,7 +321,7 @@ class TestResultRecord:
     termination_signal: ExceptionRecord, the main exception of the test.
     extra_errors: OrderedDict, all exceptions occurred during the entire
       test lifecycle. The order of occurrence is preserved.
-    result: TestResultEnum.TEAT_RESULT_*, PASS/FAIL/SKIP.
+    result: TestResultEnum.TEST_RESULT_*, PASS/FAIL/SKIP.
   """
 
   def __init__(self, t_name, t_class=None):
diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index 624f056..529035f 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -170,12 +170,19 @@ def _print_test_names(test_class):
   cls = test_class(config_parser.TestRunConfig())
   test_names = []
   try:
-    cls.setup_generated_tests()
-    test_names = cls.get_existing_test_names()
+    # Executes pre-setup procedures, this is required since it might
+    # generate test methods that we want to return as well.
+    cls._pre_run()
+    if cls.tests:
+      # Specified by run list in class.
+      test_names = list(cls.tests)
+    else:
+      # No test method specified by user, list all in test class.
+      test_names = cls.get_existing_test_names()
   except Exception:
     logging.exception('Failed to retrieve generated tests.')
   finally:
-    cls._controller_manager.unregister_controllers()
+    cls._clean_up()
   print('==========> %s <==========' % cls.TAG)
   for name in test_names:
     print(name)
diff --git a/setup.py b/setup.py
index a278b26..e99a054 100755
--- a/setup.py
+++ b/setup.py
@@ -46,13 +46,13 @@ class PyTest(test.test):
 def main():
   setuptools.setup(
       name='mobly',
-      version='1.12.1',
+      version='1.12.2',
       maintainer='Ang Li',
       maintainer_email='mobly-github@googlegroups.com',
       description='Automation framework for special end-to-end test cases',
       license='Apache2.0',
       url='https://github.com/google/mobly',
-      download_url='https://github.com/google/mobly/tarball/1.12.1',
+      download_url='https://github.com/google/mobly/tarball/1.12.2',
       packages=setuptools.find_packages(exclude=['tests']),
       include_package_data=False,
       scripts=['tools/sl4a_shell.py', 'tools/snippet_shell.py'],
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index be67437..b166357 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -510,6 +510,28 @@ class BaseTestTest(unittest.TestCase):
                         "Requested 1, Skipped 0")
     self.assertEqual(bt_cls.results.summary_str(), expected_summary)
 
+  def test_teardown_test_expects_error(self):
+
+    class MockBaseTest(base_test.BaseTestClass):
+
+      def teardown_test(self):
+        expects.expect_true(False, MSG_EXPECTED_EXCEPTION)
+
+      def test_something(self):
+        pass
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    bt_cls.run()
+    actual_record = bt_cls.results.error[0]
+    self.assertEqual(actual_record.test_name, self.mock_test_name)
+    self.assertEqual(actual_record.details, MSG_EXPECTED_EXCEPTION)
+    self.assertIsNone(actual_record.extras)
+    self.assertFalse(actual_record.extra_errors)
+    self.assertTrue(actual_record.end_time)
+    expected_summary = ("Error 1, Executed 1, Failed 0, Passed 0, "
+                        "Requested 1, Skipped 0")
+    self.assertEqual(bt_cls.results.summary_str(), expected_summary)
+
   def test_teardown_test_executed_if_test_pass(self):
     my_mock = mock.MagicMock()
 
diff --git a/tests/mobly/controllers/android_device_lib/services/logcat_test.py b/tests/mobly/controllers/android_device_lib/services/logcat_test.py
index 5c951b6..df15d76 100755
--- a/tests/mobly/controllers/android_device_lib/services/logcat_test.py
+++ b/tests/mobly/controllers/android_device_lib/services/logcat_test.py
@@ -133,6 +133,23 @@ class LogcatTest(unittest.TestCase):
               return_value=mock_android_device.MockAdbProxy('1'))
   @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
               return_value=mock_android_device.MockFastbootProxy('1'))
+  @mock.patch('mobly.utils.start_standing_subprocess')
+  @mock.patch('mobly.controllers.android_device.list_fastboot_devices',
+              return_value='1')
+  def test_start_in_fastboot_mode(self, _, start_proc_mock, FastbootProxy,
+                                  MockAdbProxy):
+    mock_serial = '1'
+    ad = android_device.AndroidDevice(serial=mock_serial)
+    logcat_service = logcat.Logcat(ad)
+    logcat_service.start()
+    # Verify start is not performed
+    self.assertFalse(logcat_service._adb_logcat_process)
+    start_proc_mock.assert_not_called()
+
+  @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+              return_value=mock_android_device.MockAdbProxy('1'))
+  @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+              return_value=mock_android_device.MockFastbootProxy('1'))
   @mock.patch('mobly.utils.create_dir')
   @mock.patch('mobly.utils.start_standing_subprocess', return_value='process')
   @mock.patch('mobly.utils.stop_standing_subprocess')
diff --git a/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py b/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py
index 162847b..16a30aa 100755
--- a/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py
+++ b/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py
@@ -15,6 +15,7 @@
 import unittest
 from unittest import mock
 
+from mobly.controllers.android_device_lib import snippet_client_v2
 from mobly.controllers.android_device_lib.services import snippet_management_service
 
 MOCK_PACKAGE = 'com.mock.package'
@@ -63,6 +64,25 @@ class SnippetManagementServiceTest(unittest.TestCase):
     mock_client.stop.assert_not_called()
 
   @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH)
+  def test_add_snippet_client_without_config(self, mock_class):
+    mock_client = mock_class.return_value
+    manager = snippet_management_service.SnippetManagementService(
+        mock.MagicMock())
+    manager.add_snippet_client('foo', MOCK_PACKAGE)
+    mock_class.assert_called_once_with(
+        package=mock.ANY, ad=mock.ANY, config=None)
+
+  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH)
+  def test_add_snippet_client_with_config(self, mock_class):
+    mock_client = mock_class.return_value
+    manager = snippet_management_service.SnippetManagementService(
+        mock.MagicMock())
+    snippet_config = snippet_client_v2.Config()
+    manager.add_snippet_client('foo', MOCK_PACKAGE, snippet_config)
+    mock_class.assert_called_once_with(
+        package=mock.ANY, ad=mock.ANY, config=snippet_config)
+
+  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH)
   def test_add_snippet_client_dup_name(self, _):
     manager = snippet_management_service.SnippetManagementService(
         mock.MagicMock())
diff --git a/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py b/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py
index 1943abb..86d889b 100644
--- a/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py
+++ b/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py
@@ -92,7 +92,9 @@ def _setup_mock_socket_file(mock_socket_create_conn, resp):
 class SnippetClientV2Test(unittest.TestCase):
   """Unit tests for SnippetClientV2."""
 
-  def _make_client(self, adb_proxy=None, mock_properties=None):
+  def _make_client(
+      self, adb_proxy=None, mock_properties=None, config=None
+  ):
     adb_proxy = adb_proxy or _MockAdbProxy(instrumented_packages=[
         (MOCK_PACKAGE_NAME, snippet_client_v2._INSTRUMENTATION_RUNNER_PACKAGE,
          MOCK_PACKAGE_NAME)
@@ -111,7 +113,9 @@ class SnippetClientV2Test(unittest.TestCase):
     }
     self.device = device
 
-    self.client = snippet_client_v2.SnippetClientV2(MOCK_PACKAGE_NAME, device)
+    self.client = snippet_client_v2.SnippetClientV2(
+        MOCK_PACKAGE_NAME, device, config
+    )
 
   def _make_client_with_extra_adb_properties(self, extra_properties):
     mock_properties = mock_android_device.DEFAULT_MOCK_PROPERTIES.copy()
@@ -408,7 +412,7 @@ class SnippetClientV2Test(unittest.TestCase):
     self.client.start_server()
     start_cmd_list = [
         'adb', 'shell',
-        (f'setsid am instrument --user {MOCK_USER_ID} -w -e action start '
+        (f'setsid am instrument --user {MOCK_USER_ID} -w -e action start  '
          f'{MOCK_SERVER_PATH}')
     ]
     self.assertListEqual(mock_start_subprocess.call_args_list,
@@ -427,7 +431,7 @@ class SnippetClientV2Test(unittest.TestCase):
     self.client.start_server()
     start_cmd_list = [
         'adb', 'shell',
-        f'setsid am instrument  -w -e action start {MOCK_SERVER_PATH}'
+        f'setsid am instrument  -w -e action start  {MOCK_SERVER_PATH}'
     ]
     self.assertListEqual(mock_start_subprocess.call_args_list,
                          [mock.call(start_cmd_list, shell=False)])
@@ -449,7 +453,7 @@ class SnippetClientV2Test(unittest.TestCase):
     self.client.start_server()
     start_cmd_list = [
         'adb', 'shell',
-        (f' am instrument --user {MOCK_USER_ID} -w -e action start '
+        (f' am instrument --user {MOCK_USER_ID} -w -e action start  '
          f'{MOCK_SERVER_PATH}')
     ]
     self.assertListEqual(mock_start_subprocess.call_args_list,
@@ -476,7 +480,7 @@ class SnippetClientV2Test(unittest.TestCase):
     self.client.start_server()
     start_cmd_list = [
         'adb', 'shell',
-        (f'nohup am instrument --user {MOCK_USER_ID} -w -e action start '
+        (f'nohup am instrument --user {MOCK_USER_ID} -w -e action start  '
          f'{MOCK_SERVER_PATH}')
     ]
     self.assertListEqual(mock_start_subprocess.call_args_list,
@@ -499,7 +503,7 @@ class SnippetClientV2Test(unittest.TestCase):
     self.client.start_server()
     start_cmd_list = [
         'adb', 'shell',
-        (f'setsid am instrument --user {MOCK_USER_ID} -w -e action start '
+        (f'setsid am instrument --user {MOCK_USER_ID} -w -e action start  '
          f'{MOCK_SERVER_PATH}')
     ]
     self.assertListEqual(mock_start_subprocess.call_args_list,
@@ -508,6 +512,28 @@ class SnippetClientV2Test(unittest.TestCase):
 
   @mock.patch('mobly.controllers.android_device_lib.snippet_client_v2.'
               'utils.start_standing_subprocess')
+  def test_start_server_with_instrument_options(self, mock_start_subprocess):
+    """Checks the starting server command with instrument options."""
+    config = snippet_client_v2.Config(
+        am_instrument_options={'key_1': 'val_1', 'key_2': 'val_2'},
+    )
+    instrument_options_str = '-e key_1 val_1 -e key_2 val_2'
+    self._make_client(config=config)
+    self._mock_server_process_starting_response(mock_start_subprocess)
+
+    self.client.start_server()
+
+    start_cmd_list = [
+        'adb', 'shell',
+        (f' am instrument --user {MOCK_USER_ID} -w -e action start '
+         f'{instrument_options_str} {MOCK_SERVER_PATH}')
+    ]
+    self.assertListEqual(mock_start_subprocess.call_args_list,
+                         [mock.call(start_cmd_list, shell=False)])
+    self.assertEqual(self.client.device_port, 1234)
+
+  @mock.patch('mobly.controllers.android_device_lib.snippet_client_v2.'
+              'utils.start_standing_subprocess')
   def test_start_server_server_crash(self, mock_start_standing_subprocess):
     """Tests that starting server process crashes."""
     self._make_client()
diff --git a/tests/mobly/controllers/android_device_test.py b/tests/mobly/controllers/android_device_test.py
index 4aa5304..3682c94 100755
--- a/tests/mobly/controllers/android_device_test.py
+++ b/tests/mobly/controllers/android_device_test.py
@@ -24,6 +24,7 @@ from mobly import runtime_test_info
 from mobly.controllers import android_device
 from mobly.controllers.android_device_lib import adb
 from mobly.controllers.android_device_lib import errors
+from mobly.controllers.android_device_lib import snippet_client_v2
 from mobly.controllers.android_device_lib.services import base_service
 from mobly.controllers.android_device_lib.services import logcat
 from tests.lib import mock_android_device
@@ -581,6 +582,16 @@ class AndroidDeviceTest(unittest.TestCase):
               return_value=mock_android_device.MockAdbProxy('1'))
   @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
               return_value=mock_android_device.MockFastbootProxy('1'))
+  @mock.patch('mobly.controllers.android_device.list_fastboot_devices',
+              return_value='1')
+  def test_AndroidDevice_is_fastboot(self, _, MockFastboot, MockAdbProxy):
+    ad = android_device.AndroidDevice(serial='1')
+    self.assertTrue(ad.is_bootloader)
+
+  @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+              return_value=mock_android_device.MockAdbProxy('1'))
+  @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+              return_value=mock_android_device.MockFastbootProxy('1'))
   @mock.patch('mobly.logger.get_log_file_timestamp')
   def test_AndroidDevice_generate_filename_default(self,
                                                    get_log_file_timestamp_mock,
@@ -1113,6 +1124,23 @@ class AndroidDeviceTest(unittest.TestCase):
   @mock.patch(
       'mobly.controllers.android_device_lib.snippet_client_v2.SnippetClientV2')
   @mock.patch('mobly.utils.get_available_host_port')
+  def test_AndroidDevice_load_snippet_with_snippet_config(
+      self, MockGetPort, MockSnippetClient, MockFastboot, MockAdbProxy):
+    ad = android_device.AndroidDevice(serial='1')
+    snippet_config = snippet_client_v2.Config()
+    ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME, snippet_config)
+    self.assertTrue(hasattr(ad, 'snippet'))
+    MockSnippetClient.assert_called_once_with(
+        package=mock.ANY, ad=mock.ANY, config=snippet_config
+    )
+
+  @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+              return_value=mock_android_device.MockAdbProxy('1'))
+  @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+              return_value=mock_android_device.MockFastbootProxy('1'))
+  @mock.patch(
+      'mobly.controllers.android_device_lib.snippet_client_v2.SnippetClientV2')
+  @mock.patch('mobly.utils.get_available_host_port')
   def test_AndroidDevice_unload_snippet(self, MockGetPort, MockSnippetClient,
                                         MockFastboot, MockAdbProxy):
     ad = android_device.AndroidDevice(serial='1')
diff --git a/tests/mobly/test_runner_test.py b/tests/mobly/test_runner_test.py
index efddc4c..0343e0d 100755
--- a/tests/mobly/test_runner_test.py
+++ b/tests/mobly/test_runner_test.py
@@ -374,20 +374,17 @@ class TestRunnerTest(unittest.TestCase):
     mock_cls_instance = mock.MagicMock()
     mock_test_class.return_value = mock_cls_instance
     test_runner._print_test_names(mock_test_class)
-    mock_cls_instance.setup_generated_tests.assert_called_once()
-    mock_cls_instance.get_existing_test_names.assert_called_once()
-    mock_cls_instance._controller_manager.unregister_controllers.assert_called_once(
-    )
+    mock_cls_instance._pre_run.assert_called_once()
+    mock_cls_instance._clean_up.assert_called_once()
 
   def test_print_test_names_with_exception(self):
     mock_test_class = mock.MagicMock()
     mock_cls_instance = mock.MagicMock()
     mock_test_class.return_value = mock_cls_instance
     test_runner._print_test_names(mock_test_class)
-    mock_cls_instance.setup_generated_tests.side_effect = Exception(
+    mock_cls_instance._pre_run.side_effect = Exception(
         'Something went wrong.')
-    mock_cls_instance._controller_manager.unregister_controllers.assert_called_once(
-    )
+    mock_cls_instance._clean_up.assert_called_once()
 
 
 if __name__ == "__main__":