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"""Test utils for UWB."""
15
16import time
17from lib.ranging_decorator import RangingTechnology
18from mobly import asserts
19from mobly.controllers import android_device
20
21WAIT_TIME_SEC = 3
22
23
24def initialize_uwb_country_code_if_necessary(ad: android_device.AndroidDevice):
25  """Sets UWB country code to US if the device does not have it set.
26
27  Note: This intentionally relies on an unstable API (shell command) since we
28  don't want to expose an API that allows users to circumvent the UWB
29  regulatory requirements.
30
31  Args:
32    ad: android device object.
33    handler: callback handler.
34  """
35  # Wait to see if UWB state is reported as enabled. If not, this could be
36  # because the country code is not set. Try forcing the country code in that
37  # case.
38  if is_technology_enabled(ad, RangingTechnology.UWB, timeout_s=60):
39    return
40
41  try:
42    ad.adb.shell(["cmd", "uwb", "force-country-code", "enabled", "US"])
43  except ad.adb.AdbError:
44    ad.log.warning("Unable to force uwb country code")
45
46  # Unable to get UWB enabled even after setting country code, abort!
47  asserts.assert_true(
48      is_technology_enabled(ad, RangingTechnology.UWB, timeout_s=60),
49      "Uwb was not enabled after setting country code",
50  )
51
52def _is_technology_state(
53    ad: android_device.AndroidDevice,
54    technology: RangingTechnology,
55    state: bool,
56    timeout_s=WAIT_TIME_SEC,
57) -> bool:
58  """Checks if the provided technology becomes enabled/disabled
59
60  Args:
61
62  ad: android device object.
63  technology: to check for enablement.
64  state: bool, True for on, False for off.
65  timeout_s: how long to wait for enablement before failing, in seconds.
66  """
67  start_time = time.time()
68  while state != ad.ranging.isTechnologyEnabled(technology):
69    if time.time() - start_time > timeout_s:
70      return False
71  return True
72
73
74def is_technology_enabled(
75    ad: android_device.AndroidDevice,
76    technology: RangingTechnology,
77    timeout_s=WAIT_TIME_SEC,
78) -> bool:
79  """Checks if the provided technology becomes enabled
80
81  Args:
82
83  ad: android device object.
84  technology: to check for enablement.
85  timeout_s: how long to wait for enablement before failing, in seconds.
86  """
87  return _is_technology_state(ad, technology, True, timeout_s)
88
89
90def set_airplane_mode(ad: android_device.AndroidDevice, state: bool):
91  """Sets the airplane mode to the given state.
92
93  Args:
94    ad: android device object.
95    state: True for Airplane mode enabled, False for disabled.
96  """
97  ad.ranging.setAirplaneMode(state)
98  start_time = time.time()
99  while get_airplane_mode(ad) != state:
100    time.sleep(0.5)
101    if time.time() - start_time > WAIT_TIME_SEC:
102      asserts.fail(f"Failed to set airplane mode to: {state}")
103
104
105def get_airplane_mode(ad: android_device.AndroidDevice) -> bool:
106  """Gets the current airplane mode setting.
107
108  Args:
109    ad: android device object.
110
111  Returns:
112    True if airplane mode On, False for Off.
113  """
114  state = ad.adb.shell(["settings", "get", "global", "airplane_mode_on"])
115  return bool(int(state.decode().strip()))
116
117def set_uwb_state_and_verify(
118    ad: android_device.AndroidDevice,
119    state: bool
120):
121  """Sets UWB state to on or off and verifies it.
122
123  Args:
124    ad: android device object.
125    state: bool, True for UWB on, False for off.
126  """
127  failure_msg = "enabled" if state else "disabled"
128  ad.uwb.setUwbEnabled(state)
129  asserts.assert_true(_is_technology_state(ad, RangingTechnology.UWB, state, timeout_s=60),
130                      "Uwb is not %s" % failure_msg)
131
132def reset_bt_state(
133    ad: android_device.AndroidDevice
134):
135  """Reset BT state to off and then on before each test.
136
137  Args:
138    ad: android device object.
139  """
140  ad.bluetooth.disableBluetooth()
141  time.sleep(3)
142  asserts.assert_false(ad.bluetooth.isBluetoothOn(), 'Bluetooth did not stop')
143  ad.bluetooth.enableBluetooth()
144  time.sleep(3)
145  asserts.assert_true(ad.bluetooth.isBluetoothOn(), 'Bluetooth did not stop')
146  # Check for BLE RSSI or BLE CS availability
147  asserts.assert_true(_is_technology_state(ad, RangingTechnology.BLE_RSSI, True, timeout_s=60),
148                      "BT is not enabled in ranging API")
149  ad.bluetooth.reset()
150
151
152def set_bt_state_and_verify(
153    ad: android_device.AndroidDevice,
154    state: bool
155):
156  """Sets BT state to on or off and verifies it.
157
158  Args:
159    ad: android device object.
160    state: bool, True for BT on, False for off.
161  """
162  failure_msg = "enabled" if state else "disabled"
163  if state:
164    ad.bluetooth.enableBluetooth()
165  else:
166    ad.bluetooth.disableBluetooth()
167  time.sleep(3)
168  asserts.assert_equal(ad.bluetooth.isBluetoothOn(), state, 'Bluetooth did not stop')
169  # Check for BLE RSSI or BLE CS availability
170  asserts.assert_true(_is_technology_state(ad, RangingTechnology.BLE_RSSI, state, timeout_s=60),
171                      "BT is not %s in ranging API" % failure_msg)
172
173
174def set_screen_rotation_landscape(
175    ad: android_device.AndroidDevice, isLandscape: bool
176):
177  """Sets screen orientation to landscape or portrait mode.
178
179  Args:
180    ad: android device object.
181    isLandscape: True for landscape mode, False for potrait.
182  """
183  ad.adb.shell(["settings", "put", "system", "accelerometer_rotation", "0"])
184  ad.adb.shell([
185      "settings",
186      "put",
187      "system",
188      "user_rotation",
189      "1" if isLandscape else "0",
190  ])
191
192
193def set_snippet_foreground_state(
194    ad: android_device.AndroidDevice, isForeground: bool
195):
196  """Sets the snippet app's foreground/background state.
197
198  Args:
199    ad: android device object.
200    isForeground: True to move snippet to foreground, False for background.
201  """
202  ad.adb.shell([
203      "cmd",
204      "uwb",
205      "simulate-app-state-change",
206      "com.google.snippet.ranging",
207      "foreground" if isForeground else "background",
208  ])
209
210
211def set_screen_state(
212    ad: android_device.AndroidDevice, on: bool
213):
214  """Sets the device screen state on/off.
215
216  Args:
217    ad: android device object.
218    on: True for screen on, False for screen off.
219  """
220  ad.adb.shell([
221      "svc", "power", "stayon", "true" if on else "false",
222  ])
223