Keep boot_completion_timeout only on handle_reboot() with default None

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):

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'):