1#  Copyright (C) 2024 The Android Open Source Project
2#
3#  Licensed under the Apache License, Version 2.0 (the "License");
4#  you may not use this file except in compliance with the License.
5#  You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#  Unless required by applicable law or agreed to in writing, software
10#  distributed under the License is distributed on an "AS IS" BASIS,
11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#  See the License for the specific language governing permissions and
13#  limitations under the License.
14
15import re
16from mobly.controllers import android_device
17from net_tests_utils.host.python import assert_utils
18
19BYTE_DECODE_UTF_8 = "utf-8"
20
21
22def set_doze_mode(ad: android_device.AndroidDevice, enable: bool) -> None:
23  if enable:
24    adb_shell(ad, "cmd battery unplug")
25    expect_dumpsys_state_with_retry(
26        ad, "deviceidle", key="mCharging", expected_state=False
27    )
28    _set_screen_state(ad, False)
29    adb_shell(ad, "dumpsys deviceidle enable deep")
30    expect_dumpsys_state_with_retry(
31        ad, "deviceidle", key="mDeepEnabled", expected_state=True
32    )
33    adb_shell(ad, "dumpsys deviceidle force-idle deep")
34    expect_dumpsys_state_with_retry(
35        ad, "deviceidle", key="mForceIdle", expected_state=True
36    )
37  else:
38    adb_shell(ad, "cmd battery reset")
39    expect_dumpsys_state_with_retry(
40        ad, "deviceidle", key="mCharging", expected_state=True
41    )
42    adb_shell(ad, "dumpsys deviceidle unforce")
43    expect_dumpsys_state_with_retry(
44        ad, "deviceidle", key="mForceIdle", expected_state=False
45    )
46
47
48def _set_screen_state(
49    ad: android_device.AndroidDevice, target_state: bool
50) -> None:
51  assert_utils.expect_with_retry(
52      predicate=lambda: _get_screen_state(ad) == target_state,
53      retry_action=lambda: adb_shell(
54          ad, "input keyevent KEYCODE_POWER"
55      ),  # Toggle power key again when retry.
56  )
57
58
59def _get_screen_state(ad: android_device.AndroidDevice) -> bool:
60  return get_value_of_key_from_dumpsys(ad, "power", "mWakefulness") == "Awake"
61
62
63def get_value_of_key_from_dumpsys(
64    ad: android_device.AndroidDevice, service: str, key: str
65) -> str:
66  output = get_dumpsys_for_service(ad, service)
67  # Search for key=value pattern from the dumpsys output.
68  # e.g. mWakefulness=Awake
69  pattern = rf"{key}=(.*)"
70  # Only look for the first occurrence.
71  match = re.search(pattern, output)
72  if match:
73    ad.log.debug(
74        "Getting key-value from dumpsys: " + key + "=" + match.group(1)
75    )
76    return match.group(1)
77  else:
78    return None
79
80
81def expect_dumpsys_state_with_retry(
82    ad: android_device.AndroidDevice,
83    service: str,
84    key: str,
85    expected_state: bool,
86    retry_interval_sec: int = 1,
87) -> None:
88  def predicate():
89    value = get_value_of_key_from_dumpsys(ad, service, key)
90    if value is None:
91      return False
92    return value.lower() == str(expected_state).lower()
93
94  assert_utils.expect_with_retry(
95      predicate=predicate,
96      retry_interval_sec=retry_interval_sec,
97  )
98
99
100def get_dumpsys_for_service(
101    ad: android_device.AndroidDevice, service: str
102) -> str:
103  return adb_shell(ad, "dumpsys " + service)
104
105
106def adb_shell(ad: android_device.AndroidDevice, shell_cmd: str) -> str:
107  """Runs adb shell command.
108
109  Args:
110    ad: Android device object.
111    shell_cmd: string of list of strings, adb shell command.
112
113  Returns:
114    string, replies from adb shell command.
115  """
116  ad.log.debug("Executing adb shell %s", shell_cmd)
117  data = ad.adb.shell(shell_cmd)
118  return data.decode(BYTE_DECODE_UTF_8).strip()
119