Optimize device loading in get_android_device_ready wrapper

PiperOrigin-RevId: 713159125
diff --git a/ui_automator/ui_automator.py b/ui_automator/ui_automator.py
index 759ddb1..f78d812 100644
--- a/ui_automator/ui_automator.py
+++ b/ui_automator/ui_automator.py
@@ -115,6 +115,9 @@
 _THIRD_PARTY_LICENSE_FILE_PATH = os.path.join(
     os.path.dirname(__file__), 'THIRD_PARTY_NOTICES.txt'
 )
+_NO_DEVICE_CONNECTED_ERROR_MESSAGE = (
+    'No Android device connected to the host computer.'
+)
 
 
 class RegTestSuiteType(enum.Enum):
@@ -220,6 +223,14 @@
     Raises:
         NoAndroidDeviceError: When no device is connected.
     """
+    # `is_adb_detectable` can detect if the device is USB connected, but it
+    # internally calls `list_adb_devices`, similar to `get_all_instances`.
+    if self._connected_device and self._connected_device.is_adb_detectable():
+      self._logger.info(
+          'Device [%s] is already connected.', self._connected_device
+      )
+      return
+
     try:
       android_devices = android_device.get_all_instances()
     except adb.AdbError as exc:
@@ -228,11 +239,10 @@
       ) from exc
 
     if not android_devices:
-      raise errors.NoAndroidDeviceError(
-          'No Android device connected to the host computer.'
-      )
+      raise errors.NoAndroidDeviceError(_NO_DEVICE_CONNECTED_ERROR_MESSAGE)
+
     self._connected_device = android_devices[0]
-    self._logger.info(f'connected device: [{self._connected_device}]')
+    self._logger.info('connected device: [%s]', self._connected_device)
 
   def load_snippet(self):
     """Loads needed mobly snippet installed on android device.
diff --git a/ui_automator/ui_automator_test.py b/ui_automator/ui_automator_test.py
index 383f629..8c27e91 100644
--- a/ui_automator/ui_automator_test.py
+++ b/ui_automator/ui_automator_test.py
@@ -108,7 +108,8 @@
       ui_automator.run()
 
     mock_exit.assert_called_once()
-    self.assertIn("""Copyright 2024 Google LLC
+    self.assertIn(
+        """Copyright 2024 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -120,7 +121,9 @@
 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.""", mock_print.call_args.args[0])
+limitations under the License.""",
+        mock_print.call_args.args[0],
+    )
 
   @mock.patch.object(android_device, 'get_all_instances', autospec=True)
   def test_load_device_raises_an_error_when_adb_not_installed(
@@ -155,6 +158,63 @@
       self.ui_automator.load_device()
     self.assertEqual(cm.output, ['INFO:root:connected device: [1]'])
 
+  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
+  def test_load_device_should_raise_no_device_error_with_disconnected_device(
+      self, mock_get_all_instances
+  ):
+    mock_get_all_instances.return_value = [self.mock_android_device]
+    self.ui_automator.load_device()
+    self.mock_android_device.is_adb_detectable.return_value = False
+    mock_get_all_instances.return_value = []
+
+    with self.assertRaisesRegex(
+        errors.NoAndroidDeviceError,
+        r'No Android device connected to the host computer\.',
+    ):
+      self.ui_automator.load_device()
+
+    self.assertEqual(mock_get_all_instances.call_count, 2)
+
+  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
+  def test_load_device_should_connect_to_the_device_on_reconnection(
+      self, mock_get_all_instances
+  ):
+    mock_get_all_instances.return_value = [self.mock_android_device, 2]
+    self.ui_automator.load_device()
+    self.mock_android_device.is_adb_detectable.return_value = False
+
+    with self.assertLogs() as cm:
+      self.ui_automator.load_device()
+
+    self.assertEqual(mock_get_all_instances.call_count, 2)
+    self.assertEqual(
+        cm.output,
+        [f'INFO:root:connected device: [{self.mock_android_device}]'],
+    )
+
+  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
+  def test_load_device_should_not_call_get_all_instances_with_connected_device(
+      self, mock_get_all_instances
+  ):
+    mock_get_all_instances.return_value = [self.mock_android_device, 2]
+    self.mock_android_device.is_adb_detectable.return_value = True
+
+    with self.assertLogs() as cm:
+      self.ui_automator.load_device()
+      self.ui_automator.load_device()
+
+    mock_get_all_instances.assert_called_once()
+    self.assertEqual(
+        cm.output,
+        [
+            f'INFO:root:connected device: [{self.mock_android_device}]',
+            (
+                f'INFO:root:Device [{self.mock_android_device}] is already'
+                ' connected.'
+            ),
+        ],
+    )
+
   def test_load_snippet_without_load_device_raises_error(self):
     with self.assertRaises(errors.NoAndroidDeviceError):
       self.ui_automator.load_snippet()