From cfc8a85432978cfb5ff8b605f9ddf4c5ee57b03f Mon Sep 17 00:00:00 2001
From: Jimi Chen <jimictw@google.com>
Date: Wed, 28 May 2025 13:13:17 +0000
Subject: [PATCH] Force device unroot in the multi-device setup

In mobly, it will root the device if the device is rootable.
If there is some operation needs root permission,
the test result will be inconsistent between user and userdebug build.

Flag: EXEMPT mainline
Bug: 419423721
Test: atest CtsConnectivityMultiDevicesTestCases
Change-Id: Ib962b565aeaf8f0e1737b700a16420277e088e6b
---
 staticlibs/testutils/host/python/adb_utils.py | 72 +++++++++++++++++++
 .../host/python/multi_devices_test_base.py    | 15 ++++
 2 files changed, 87 insertions(+)

diff --git a/staticlibs/testutils/host/python/adb_utils.py b/staticlibs/testutils/host/python/adb_utils.py
index 023cb01164..e3260a9920 100644
--- a/staticlibs/testutils/host/python/adb_utils.py
+++ b/staticlibs/testutils/host/python/adb_utils.py
@@ -12,8 +12,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import logging
 import re
+import time
+import types
+from mobly import utils
 from mobly.controllers import android_device
+from mobly.controllers.android_device_lib import adb
 from net_tests_utils.host.python import assert_utils
 
 BYTE_DECODE_UTF_8 = 'utf-8'
@@ -103,6 +108,73 @@ def get_dumpsys_for_service(
   return adb_shell(ad, 'dumpsys ' + service)
 
 
+def unroot_adb(self) -> bytes:
+  """Executes the 'adb unroot' command on the device with retry logic.
+
+  This method is intended to be dynamically attached to an AdbProxy instance
+  (or similar object) that provides an '_exec_adb_cmd' method for executing
+  ADB commands. It attempts to unroot the device, and if an AdbError occurs,
+  it retries the command a specified number of times with an exponentially
+  increasing delay.
+
+  Args:
+    self: The instance of the ADB client (e.g., an AdbProxy object) that has the
+      '_exec_adb_cmd' method available to run ADB commands.
+
+  Returns:
+    bytes: The raw byte output from the 'adb unroot' command's successful
+      execution.
+
+  Raises:
+    adb.AdbError: If the 'adb unroot' command fails after all specified
+      retry attempts.
+  """
+  retry_interval = adb.ADB_ROOT_RETRY_ATTEMPT_INTERVAL_SEC
+  for attempt in range(adb.ADB_ROOT_RETRY_ATTEMPTS):
+    try:
+      # Assuming 'self' here is an object that has a '_exec_adb_cmd' method,
+      # similar to your original 'root' method's context (e.g., an AdbProxy).
+      return self._exec_adb_cmd(
+          'unroot', args=None, shell=False, timeout=None, stderr=None
+      )
+    except adb.AdbError as e:
+      if attempt + 1 < adb.ADB_ROOT_RETRY_ATTEMPTS:
+        # Decode stderr for logging, as Mobly often provides bytes
+        decoded_stderr = e.stderr.decode('utf-8', errors='ignore').strip()
+        logging.error(
+            f'Retry the command "{utils.cli_cmd_to_string(e.cmd)}" since Error'
+            f' "{decoded_stderr}" occurred (attempt'
+            f' {attempt + 1}/{adb.ADB_ROOT_RETRY_ATTEMPTS}).'
+        )
+        time.sleep(retry_interval)
+        retry_interval *= 2
+      else:
+        raise e
+
+
+def unroot(ad: android_device.AndroidDevice):
+  """Dynamically unroots an Android device by attaching 'unroot_adb' to its ADB proxy.
+
+  This function addresses a potential issue where Mobly might attempt to root
+  a device upon registration. When running user or userdebug builds,
+  this automatic rooting can lead to unexpected or inconsistent device behavior.
+  To mitigate this, this function temporarily attaches the 'unroot_adb' method
+  to the device's ADB proxy (`ad.adb`), executes it, and then removes it,
+  ensuring the device returns to an unrooted state and the ADB proxy's state
+  remains clean.
+
+  Args:
+    ad: The `android_device.AndroidDevice` instance to be unrooted.
+  """
+  attach_method_name = 'unroot_adb'
+  try:
+    setattr(ad.adb, attach_method_name, types.MethodType(unroot_adb, ad.adb))
+    getattr(ad.adb, attach_method_name)()
+    ad.adb.wait_for_device()
+  finally:
+    delattr(ad.adb, attach_method_name)
+
+
 def adb_shell(ad: android_device.AndroidDevice, shell_cmd: str) -> str:
   """Runs adb shell command.
 
diff --git a/staticlibs/testutils/host/python/multi_devices_test_base.py b/staticlibs/testutils/host/python/multi_devices_test_base.py
index 6db9577493..a0b15ccd98 100644
--- a/staticlibs/testutils/host/python/multi_devices_test_base.py
+++ b/staticlibs/testutils/host/python/multi_devices_test_base.py
@@ -24,9 +24,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from mobly import asserts
 from mobly import base_test
 from mobly import utils
 from mobly.controllers import android_device
+from net_tests_utils.host.python import adb_utils
 
 CONNECTIVITY_MULTI_DEVICES_SNIPPET_PACKAGE = 'com.google.snippet.connectivity'
 
@@ -39,6 +41,19 @@ class MultiDevicesTestBase(base_test.BaseTestClass):
         android_device, min_number=2
     )
 
+    if self.clientDevice.is_adb_root:
+      adb_utils.unroot(self.clientDevice)
+
+    if self.serverDevice.is_adb_root:
+      adb_utils.unroot(self.serverDevice)
+
+    asserts.assert_false(
+        self.clientDevice.is_adb_root, 'client device should be unroot'
+    )
+    asserts.assert_false(
+        self.serverDevice.is_adb_root, 'server device should be unroot'
+    )
+
     def setup_device(device):
       device.load_snippet(
           'connectivity_multi_devices_snippet',