diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py index 06748d76..361b541f 100644 --- a/mobly/controllers/android_device.py +++ b/mobly/controllers/android_device.py @@ -709,9 +709,7 @@ def update_serial(self, new_serial): self.fastboot.serial = new_serial @contextlib.contextmanager - def handle_reboot( - self, boot_completion_timeout=DEFAULT_TIMEOUT_BOOT_COMPLETION_SECOND - ): + def handle_reboot(self, boot_completion_timeout=None): """Properly manage the service life cycle when the device needs to temporarily disconnect. @@ -723,7 +721,8 @@ def handle_reboot( Args: boot_completion_timeout: float, the number of seconds to wait for the - device to finish booting after the context exits. This only bounds the + device to finish booting after the context exits. If None, uses the + default timeout of `wait_for_boot_completion`. This only bounds the boot completion wait, not the code executed inside the context. """ live_service_names = self.services.list_live_services() @@ -739,7 +738,10 @@ def handle_reboot( try: yield finally: - self.wait_for_boot_completion(timeout=boot_completion_timeout) + if boot_completion_timeout is None: + self.wait_for_boot_completion() + else: + self.wait_for_boot_completion(timeout=boot_completion_timeout) # On boot completion, invalidate the `build_info` cache since any # value it had from before boot completion is potentially invalid. # If the value gets set after the final invalidation and before @@ -1191,9 +1193,7 @@ def is_adb_detectable(self): return True return False - def reboot( - self, boot_completion_timeout=DEFAULT_TIMEOUT_BOOT_COMPLETION_SECOND - ): + def reboot(self): """Reboots the device. Generally one should use this method to reboot the device instead of @@ -1203,19 +1203,13 @@ def reboot( This method is blocking and only returns when the reboot has completed and the services restored. - Args: - boot_completion_timeout: float, the number of seconds to wait for the - device to finish booting. This has no effect when the device is in - bootloader mode, since this method does not wait for boot completion - in that case. - Raises: Error: Waiting for completion timed out. """ if self.is_bootloader: self.fastboot.reboot() return - with self.handle_reboot(boot_completion_timeout=boot_completion_timeout): + with self.handle_reboot(): self.adb.reboot() def __getattr__(self, name): diff --git a/tests/mobly/controllers/android_device_test.py b/tests/mobly/controllers/android_device_test.py index 940906c5..28a5185b 100755 --- a/tests/mobly/controllers/android_device_test.py +++ b/tests/mobly/controllers/android_device_test.py @@ -2151,24 +2151,6 @@ def test_AndroidDevice_handle_reboot_custom_timeout( pass mock_wait.assert_called_once_with(timeout=120) - @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'), - ) - def test_AndroidDevice_reboot_custom_timeout( - self, - FastbootProxy, - MockAdbProxy, - ): - ad = android_device.AndroidDevice(serial='1') - with mock.patch.object(ad, 'wait_for_boot_completion') as mock_wait: - ad.reboot(boot_completion_timeout=60) - mock_wait.assert_called_once_with(timeout=60) - def test_AndroidDevice_getattr_without_services(self): ad = android_device.AndroidDevice.__new__(android_device.AndroidDevice) with self.assertRaisesRegex(AttributeError, 'services'):