Add adb_utils for multidevice test

This commit adds the necessary code for 'adb_utils' to enable APF
multi-device testing which will be introduced in follow-up patches.
This includes utilities of issuing shell command, parsing dumpsys
key-value pairs, controlling doze mode and screen state.

Additionally, unit tests are added to verify the basic functionality.

Test: atest NetworkStaticLibHostPythonTests

Bug: 335368434
Change-Id: I580e91b494580ec9105c8988972845d38c104031
diff --git a/staticlibs/tests/unit/host/python/adb_utils_test.py b/staticlibs/tests/unit/host/python/adb_utils_test.py
new file mode 100644
index 0000000..b75d464
--- /dev/null
+++ b/staticlibs/tests/unit/host/python/adb_utils_test.py
@@ -0,0 +1,126 @@
+#  Copyright (C) 2024 The Android Open Source Project
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+
+import unittest
+from unittest.mock import MagicMock, patch
+from net_tests_utils.host.python import adb_utils
+from net_tests_utils.host.python.assert_utils import UnexpectedBehaviorError
+
+
+class TestAdbUtils(unittest.TestCase):
+
+  def setUp(self):
+    self.mock_ad = MagicMock()  # Mock Android device object
+    self.mock_ad.log = MagicMock()
+    self.mock_ad.adb.shell.return_value = b""  # Default empty return for shell
+
+  @patch(
+      "net_tests_utils.host.python.adb_utils.expect_dumpsys_state_with_retry"
+  )
+  @patch("net_tests_utils.host.python.adb_utils._set_screen_state")
+  def test_set_doze_mode_enable(
+      self, mock_set_screen_state, mock_expect_dumpsys_state
+  ):
+    adb_utils.set_doze_mode(self.mock_ad, True)
+    mock_set_screen_state.assert_called_once_with(self.mock_ad, False)
+
+  @patch(
+      "net_tests_utils.host.python.adb_utils.expect_dumpsys_state_with_retry"
+  )
+  def test_set_doze_mode_disable(self, mock_expect_dumpsys_state):
+    adb_utils.set_doze_mode(self.mock_ad, False)
+
+  @patch("net_tests_utils.host.python.adb_utils._get_screen_state")
+  def test_set_screen_state_success(self, mock_get_screen_state):
+    mock_get_screen_state.side_effect = [False, True]  # Simulate toggle
+    adb_utils._set_screen_state(self.mock_ad, True)
+
+  @patch("net_tests_utils.host.python.adb_utils._get_screen_state")
+  def test_set_screen_state_failure(self, mock_get_screen_state):
+    mock_get_screen_state.return_value = False  # State doesn't change
+    with self.assertRaises(UnexpectedBehaviorError):
+      adb_utils._set_screen_state(self.mock_ad, True)
+
+  @patch("net_tests_utils.host.python.adb_utils.get_value_of_key_from_dumpsys")
+  def test_get_screen_state(self, mock_get_value):
+    # Test cases for different return values of get_value_of_key_from_dumpsys
+    # TODO: Make it parameterized.
+    test_cases = [
+        ("Awake", True),
+        ("Asleep", False),
+        ("Dozing", False),
+        ("SomeOtherState", False),
+    ]
+
+    for state_str, expected_result in test_cases:
+      mock_get_value.return_value = state_str
+      self.assertEqual(
+          adb_utils._get_screen_state(self.mock_ad), expected_result
+      )
+
+  def test_get_value_of_key_from_dumpsys(self):
+    self.mock_ad.adb.shell.return_value = (
+        b"mWakefulness=Awake\nmOtherKey=SomeValue"
+    )
+    result = adb_utils.get_value_of_key_from_dumpsys(
+        self.mock_ad, "power", "mWakefulness"
+    )
+    self.assertEqual(result, "Awake")
+
+  @patch("net_tests_utils.host.python.adb_utils.get_value_of_key_from_dumpsys")
+  def test_expect_dumpsys_state_with_retry_success(self, mock_get_value):
+    # Test cases for different combinations of expected_state and get_value return value
+    # TODO: Make it parameterized.
+    test_cases = [
+        (True, ["true"]),  # Expect True, get True
+        (False, ["false"]),  # Expect False, get False
+        (
+            True,
+            ["false", "true"],
+        ),  # Expect True, get False which is unexpected, then get True
+        (
+            False,
+            ["true", "false"],
+        ),  # Expect False, get True which is unexpected, then get False
+    ]
+
+    for expected_state, returned_value in test_cases:
+      mock_get_value.side_effect = returned_value
+      # Verify the method returns and does not throw.
+      adb_utils.expect_dumpsys_state_with_retry(
+          self.mock_ad, "service", "key", expected_state, 0
+      )
+
+  @patch("net_tests_utils.host.python.adb_utils.get_value_of_key_from_dumpsys")
+  def test_expect_dumpsys_state_with_retry_failure(self, mock_get_value):
+    mock_get_value.return_value = "false"
+    with self.assertRaises(UnexpectedBehaviorError):
+      adb_utils.expect_dumpsys_state_with_retry(
+          self.mock_ad, "service", "key", True, 0
+      )
+
+  @patch("net_tests_utils.host.python.adb_utils.get_value_of_key_from_dumpsys")
+  def test_expect_dumpsys_state_with_retry_not_found(self, mock_get_value):
+    # Simulate the get_value_of_key_from_dumpsys cannot find the give key.
+    mock_get_value.return_value = None
+
+    # Expect the function to raise UnexpectedBehaviorError due to the exception
+    with self.assertRaises(UnexpectedBehaviorError):
+      adb_utils.expect_dumpsys_state_with_retry(
+          self.mock_ad, "service", "key", True
+      )
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/staticlibs/tests/unit/host/python/run_tests.py b/staticlibs/tests/unit/host/python/run_tests.py
index 8059568..7d58c6e 100644
--- a/staticlibs/tests/unit/host/python/run_tests.py
+++ b/staticlibs/tests/unit/host/python/run_tests.py
@@ -18,6 +18,7 @@
 
 # Import all unittest classes here, so it can be discovered by unittest module.
 # TODO: make the tests can be executed without manually import classes.
+from host.python.adb_utils_test import TestAdbUtils
 from host.python.assert_utils_test import TestAssertUtils
 
 
diff --git a/staticlibs/testutils/host/python/adb_utils.py b/staticlibs/testutils/host/python/adb_utils.py
new file mode 100644
index 0000000..13c0646
--- /dev/null
+++ b/staticlibs/testutils/host/python/adb_utils.py
@@ -0,0 +1,118 @@
+#  Copyright (C) 2024 The Android Open Source Project
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+
+import re
+from mobly.controllers import android_device
+from net_tests_utils.host.python import assert_utils
+
+BYTE_DECODE_UTF_8 = "utf-8"
+
+
+def set_doze_mode(ad: android_device.AndroidDevice, enable: bool) -> None:
+  if enable:
+    adb_shell(ad, "cmd battery unplug")
+    expect_dumpsys_state_with_retry(
+        ad, "deviceidle", key="mCharging", expected_state=False
+    )
+    _set_screen_state(ad, False)
+    adb_shell(ad, "dumpsys deviceidle enable deep")
+    expect_dumpsys_state_with_retry(
+        ad, "deviceidle", key="mDeepEnabled", expected_state=True
+    )
+    adb_shell(ad, "dumpsys deviceidle force-idle deep")
+    expect_dumpsys_state_with_retry(
+        ad, "deviceidle", key="mForceIdle", expected_state=True
+    )
+  else:
+    adb_shell(ad, "cmd battery reset")
+    expect_dumpsys_state_with_retry(
+        ad, "deviceidle", key="mCharging", expected_state=True
+    )
+    adb_shell(ad, "dumpsys deviceidle unforce")
+    expect_dumpsys_state_with_retry(
+        ad, "deviceidle", key="mForceIdle", expected_state=False
+    )
+
+
+def _set_screen_state(
+    ad: android_device.AndroidDevice, target_state: bool
+) -> None:
+  assert_utils.expect_with_retry(
+      predicate=lambda: _get_screen_state(ad) == target_state,
+      retry_action=lambda: adb_shell(
+          ad, "input keyevent KEYCODE_POWER"
+      ),  # Toggle power key again when retry.
+  )
+
+
+def _get_screen_state(ad: android_device.AndroidDevice) -> bool:
+  return get_value_of_key_from_dumpsys(ad, "power", "mWakefulness") == "Awake"
+
+
+def get_value_of_key_from_dumpsys(
+    ad: android_device.AndroidDevice, service: str, key: str
+) -> str:
+  output = get_dumpsys_for_service(ad, service)
+  # Search for key=value pattern from the dumpsys output.
+  # e.g. mWakefulness=Awake
+  pattern = rf"{key}=(.*)"
+  # Only look for the first occurrence.
+  match = re.search(pattern, output)
+  if match:
+    ad.log.debug(
+        "Getting key-value from dumpsys: " + key + "=" + match.group(1)
+    )
+    return match.group(1)
+  else:
+    return None
+
+
+def expect_dumpsys_state_with_retry(
+    ad: android_device.AndroidDevice,
+    service: str,
+    key: str,
+    expected_state: bool,
+    retry_interval_sec: int = 1,
+) -> None:
+  def predicate():
+    value = get_value_of_key_from_dumpsys(ad, service, key)
+    if value is None:
+      return False
+    return value.lower() == str(expected_state).lower()
+
+  assert_utils.expect_with_retry(
+      predicate=predicate,
+      retry_interval_sec=retry_interval_sec,
+  )
+
+
+def get_dumpsys_for_service(
+    ad: android_device.AndroidDevice, service: str
+) -> str:
+  return adb_shell(ad, "dumpsys " + service)
+
+
+def adb_shell(ad: android_device.AndroidDevice, shell_cmd: str) -> str:
+  """Runs adb shell command.
+
+  Args:
+    ad: Android device object.
+    shell_cmd: string of list of strings, adb shell command.
+
+  Returns:
+    string, replies from adb shell command.
+  """
+  ad.log.debug("Executing adb shell %s", shell_cmd)
+  data = ad.adb.shell(shell_cmd)
+  return data.decode(BYTE_DECODE_UTF_8).strip()