diff --git a/mobly/controllers/android_device_lib/adb.py b/mobly/controllers/android_device_lib/adb.py index 006758ee..29faec7c 100644 --- a/mobly/controllers/android_device_lib/adb.py +++ b/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 diff --git a/tests/mobly/controllers/android_device_lib/adb_test.py b/tests/mobly/controllers/android_device_lib/adb_test.py index 0f23e126..6da423ed 100755 --- a/tests/mobly/controllers/android_device_lib/adb_test.py +++ b/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