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 logging
17import random
18import time
19from typing import List, Optional
20from lib import generic_ranging_decorator
21from mobly import asserts
22from mobly.controllers import android_device
23from mobly.controllers.android_device_lib import adb
24from mobly.controllers.android_device_lib import callback_handler_v2
25
26WAIT_TIME_SEC = 3
27
28
29def verify_uwb_state_callback(
30        ad: android_device.AndroidDevice,
31        uwb_event: str,
32        handler: Optional[callback_handler_v2.CallbackHandlerV2] = None,
33        timeout: int = WAIT_TIME_SEC,
34) -> bool:
35    """Verifies expected UWB callback is received.
36
37    Args:
38      ad: android device object.
39      uwb_event: expected callback event.
40      handler: callback handler.
41      timeout: timeout for callback event.
42
43    Returns:
44      True if expected callback is received, False if not.
45    """
46    callback_status = False
47    callback_key = None
48    start_time = time.time()
49    if handler is None:
50        callback_key = "uwb_state_%s" % random.randint(1, 100)
51        handler = ad.uwb.registerUwbAdapterStateCallback(callback_key)
52    # wait until expected callback is received.
53    while time.time() - start_time < timeout and not callback_status:
54        time.sleep(0.1)
55        events = handler.getAll("UwbAdapterStateCallback")
56        for event in events:
57            event_received = event.data["uwbAdapterStateEvent"]
58            logging.debug("Received event - %s", event_received)
59            if event_received == uwb_event:
60                logging.debug("Received the '%s' callback in %ss", uwb_event,
61                              round(time.time() - start_time, 2))
62                callback_status = True
63                break
64    if callback_key is not None:
65        ad.uwb.unregisterUwbAdapterStateCallback(callback_key)
66    return callback_status
67
68
69def get_uwb_state(ad: android_device.AndroidDevice) -> bool:
70    """Gets the current UWB state.
71
72    Args:
73      ad: android device object.
74
75    Returns:
76      UWB state, True if enabled, False if not.
77    """
78    if ad.build_info["build_id"].startswith("S"):
79        uwb_state = bool(ad.uwb.getAdapterState())
80    else:
81        uwb_state = ad.uwb.isUwbEnabled()
82    return uwb_state
83
84
85def set_uwb_state_and_verify(
86        ad: android_device.AndroidDevice,
87        state: bool,
88        handler: Optional[callback_handler_v2.CallbackHandlerV2] = None,
89):
90    """Sets UWB state to on or off and verifies it.
91
92    Args:
93      ad: android device object.
94      state: bool, True for UWB on, False for off.
95      handler: callback_handler.
96    """
97    failure_msg = "enabled" if state else "disabled"
98    ad.uwb.setUwbEnabled(state)
99    event_str = "Inactive" if state else "Disabled"
100    asserts.assert_true(verify_uwb_state_callback(ad, event_str, handler),
101                        "Uwb is not %s" % failure_msg)
102
103
104def verify_peer_found(ranging_dut: generic_ranging_decorator.GenericRangingDecorator,
105                      peer_addr: List[int], session: int = 0):
106    """Verifies if the UWB peer is found.
107
108    Args:
109      ranging_dut: uwb ranging device.
110      peer_addr: uwb peer device address.
111      session: session id.
112    """
113    ranging_dut.ad.log.info("Look for peer: %s" % peer_addr)
114    start_time = time.time()
115    while not ranging_dut.is_uwb_peer_found(peer_addr, session):
116        if time.time() - start_time > WAIT_TIME_SEC:
117            asserts.fail("UWB peer with address %s not found" % peer_addr)
118    logging.info("Peer %s found in %s seconds", peer_addr,
119                 round(time.time() - start_time, 2))
120
121
122def set_airplane_mode(ad: android_device.AndroidDevice, state: bool):
123    """Sets the airplane mode to the given state.
124
125    Args:
126      ad: android device object.
127      state: bool, True for Airplane mode on, False for off.
128    """
129    ad.uwb.setAirplaneMode(state)
130    start_time = time.time()
131    while get_airplane_mode(ad) != state:
132        time.sleep(0.5)
133        if time.time() - start_time > WAIT_TIME_SEC:
134            asserts.fail("Failed to set airplane mode to: %s" % state)
135
136
137def get_airplane_mode(ad: android_device.AndroidDevice) -> bool:
138    """Gets the airplane mode.
139
140    Args:
141      ad: android device object.
142
143    Returns:
144      True if airplane mode On, False for Off.
145    """
146    state = ad.adb.shell(["settings", "get", "global", "airplane_mode_on"])
147    return bool(int(state.decode().strip()))
148
149
150def set_screen_rotation(ad: android_device.AndroidDevice, val: int):
151    """Sets screen orientation to landscape or portrait mode.
152
153    Args:
154      ad: android device object.
155      val: False for potrait, True 1 for landscape mode.
156    """
157    ad.adb.shell(["settings", "put", "system", "accelerometer_rotation", "0"])
158    ad.adb.shell(["settings", "put", "system", "user_rotation", str(val)])
159
160
161def initialize_uwb_country_code_if_not_set(
162        ad: android_device.AndroidDevice,
163        handler: Optional[callback_handler_v2.CallbackHandlerV2] = None,
164):
165    """Sets UWB country code to US if the device does not have it set.
166
167    Note: This intentionally relies on an unstable API (shell command) since we
168    don't want to expose an API that allows users to circumvent the UWB
169    regulatory requirements.
170
171    Args:
172      ad: android device object.
173      handler: callback handler.
174    """
175    # Wait to see if UWB state is reported as enabled. If not, this could be
176    # because the country code is not set. Try forcing the country code in that
177    # case.
178    state = verify_uwb_state_callback(
179        ad=ad, uwb_event="Inactive", handler=handler, timeout=120
180    )
181
182    # Country code already available, nothing to do.
183    if state:
184        return
185    try:
186        ad.adb.shell(["cmd", "uwb", "force-country-code", "enabled", "US"])
187    except adb.AdbError:
188        logging.warning("Unable to force country code")
189
190    # Unable to get UWB enabled even after setting country code, abort!
191    asserts.fail(
192        not verify_uwb_state_callback(
193            ad=ad, uwb_event="Inactive", handler=handler, timeout=120
194        ),
195        "Uwb is not enabled",
196    )
197