Update adb connect method to disconnect first

mobly/controllers/android_device_lib/adb.py

@@ -354,9 +354,24 @@ def current_user_id(self) -> int:
     # Multi-user is not supported in SDK < 21, only user 0 exists.
     return 0
 
+  def disconnect(self, address) -> bytes:
+    """Executes the `adb disconnect` command.
+
+    Args:
+      address: string, the address of the Android instance to disconnect from.
+
+    Returns:
+      The stdout content.
+    """
+    return self._exec_adb_cmd(
+        'disconnect', address, shell=False, timeout=None, stderr=None
+    )
+
   def connect(self, address) -> bytes:
     """Executes the `adb connect` command with proper status checking.
 
+    Disconnects first to flush any stale or zombie socket states on the ADB server.
+
     Args:
       address: string, the address of the Android instance to connect to.
 
@@ -366,6 +381,12 @@ def connect(self, address) -> bytes:
     Raises:
       AdbError: if the connection failed.
     """
+    # Flush any stale/zombie socket states before connecting
+    try:
+      self.disconnect(address)
+    except AdbError:
+      pass
+
     stdout = self._exec_adb_cmd(
         'connect', address, shell=False, timeout=None, stderr=None
     )
@@ -561,6 +582,18 @@ def root(self) -> bytes:
 
           # Buffer between "adb root" commands.
           time.sleep(retry_interval)
+          if ':' in self.serial:
+            try:
+              logging.debug(
+                  'Reconnecting TCP device %s before root retry...', self.serial
+              )
+              self.connect(self.serial)
+            except AdbError as connect_err:
+              logging.debug(
+                  'Reconnecting to %s during root retry failed: %s',
+                  self.serial,
+                  connect_err,
+              )
           retry_interval *= 2
         else:
           raise e

tests/mobly/controllers/android_device_lib/adb_test.py

@@ -560,6 +560,17 @@ def test_connect_fail_error_fields_are_bytes(self, mock_run_command):
     self.assertIsInstance(cm.exception.stdout, bytes)
     self.assertIsInstance(cm.exception.stderr, bytes)
 
+  @mock.patch('mobly.utils.run_command')
+  def test_connect_ignores_disconnect_error(self, mock_run_command):
+    mock_address = 'localhost:1234'
+    mock_run_command.side_effect = [
+        (1, b'', b'error: no such device'),
+        (0, f'connected to {mock_address}'.encode('utf-8'), b''),
+    ]
+    out = adb.AdbProxy().connect(mock_address)
+    self.assertEqual('connected to localhost:1234', out.decode('utf-8'))
+    self.assertEqual(mock_run_command.call_count, 2)
+
   def test_getprop(self):
     with mock.patch.object(adb.AdbProxy, '_exec_cmd') as mock_exec_cmd:
       mock_exec_cmd.return_value = b'blah'
@@ -936,6 +947,29 @@ def test_root_success_with_retry_after_timeout(
     self.assertEqual(mock_sleep.call_count, 1)
     mock_sleep.assert_called_with(10)
 
+  @mock.patch('time.sleep', return_value=mock.MagicMock())
+  @mock.patch.object(adb.AdbProxy, 'connect')
+  @mock.patch.object(adb.AdbProxy, '_exec_cmd')
+  def test_root_tcp_reconnect_on_retry(
+      self, mock_exec_cmd, mock_connect, mock_sleep
+  ):
+    tcp_serial = 'localhost:1234'
+    mock_exec_cmd.side_effect = [
+        adb.AdbError('adb root', '', MOCK_ROOT_ERROR_OUTPUT, 1),
+        adb.AdbError('adb root', '', MOCK_ROOT_ERROR_OUTPUT, 1),
+        MOCK_ROOT_SUCCESS_OUTPUT,
+    ]
+    mock_connect.side_effect = [
+        adb.AdbError('connect localhost:1234', b'', b'', 0),
+        b'connected to localhost:1234',
+    ]
+    output = adb.AdbProxy(tcp_serial).root()
+    self.assertEqual(output, MOCK_ROOT_SUCCESS_OUTPUT)
+    mock_connect.assert_has_calls(
+        [mock.call(tcp_serial), mock.call(tcp_serial)]
+    )
+    self.assertEqual(mock_sleep.call_count, 2)
+
   def test_has_shell_command_called_correctly(self):
     with mock.patch.object(adb.AdbProxy, '_exec_cmd') as mock_exec_cmd:
       mock_exec_cmd.return_value = MOCK_DEFAULT_COMMAND_OUTPUT