1"""Tests for UwbManager APIs."""
2
3import random
4import sys
5from lib import uwb_base_test
6from mobly import asserts
7from mobly import config_parser
8from mobly import suite_runner
9from mobly.controllers import android_device
10from mobly.controllers.android_device_lib import callback_handler_v2
11from test_utils import uwb_test_utils
12
13_TEST_CASES = (
14    "test_toggle_uwb_state",
15    "test_uwb_state_with_airplane_mode_toggle",
16    "test_uwb_state_off_with_airplane_mode_toggle",
17)
18
19
20class UwbManagerTest(uwb_base_test.UwbBaseTest):
21  """Tests for UwbManager platform APIs.
22
23  Attributes:
24    android_devices: list of android device objects.
25  """
26
27  def __init__(self, configs: config_parser.TestRunConfig):
28    """Init method for the test class.
29
30    Args:
31      configs: A config_parser.TestRunConfig object.
32    """
33    super().__init__(configs)
34    self.tests = _TEST_CASES
35
36  def setup_class(self):
37    super().setup_class()
38    self.dut = self.android_devices[0]
39
40  def setup_test(self):
41    super().setup_test()
42    self.callback = "uwb_state_%s" % random.randint(1, 100)
43    self.handler = self.dut.uwb.registerUwbAdapterStateCallback(self.callback)
44    if not uwb_test_utils.get_uwb_state(self.dut):
45      uwb_test_utils.set_uwb_state_and_verify(self.dut, True, self.handler)
46
47  def teardown_test(self):
48    super().teardown_test()
49    uwb_test_utils.set_airplane_mode(self.dut, False)
50    self.dut.uwb.unregisterUwbAdapterStateCallback(self.callback)
51
52  ### Helper methods ###
53
54  def _test_uwb_state_after_reboot(
55      self,
56      dut: android_device.AndroidDevice,
57      state: bool,
58      handler: callback_handler_v2.CallbackHandlerV2,
59  ):
60    """Sets UWB state and verifies it is persistent after reboot.
61
62    Args:
63      dut: android device object.
64      state: bool, True for UWB mode on, False for off.
65      handler: callback handler.
66    """
67    uwb_test_utils.set_uwb_state_and_verify(dut, state)
68    dut.reboot()
69    uwb_test_utils.initialize_uwb_country_code_if_not_set(dut)
70    state_after_reboot = uwb_test_utils.get_uwb_state(dut)
71    asserts.assert_equal(
72        state, state_after_reboot,
73        "Uwb state before reboot: %s;  after reboot: %s" %
74        (state, state_after_reboot))
75
76  def _verify_uwb_state_with_airplane_mode(
77      self,
78      dut: android_device.AndroidDevice,
79      prev_uwb_state: bool,
80      expected_uwb_state: bool,
81      event_str: str,
82      handler: callback_handler_v2.CallbackHandlerV2,
83  ):
84    """Verifies UWB state with airplane mode.
85
86    Args:
87      dut: android device object.
88      prev_uwb_state: previous UWB state.
89      expected_uwb_state: expected UWB state.
90      event_str: callback event string.
91      handler: callback handler.
92
93    Returns:
94      True if test results in expected UWB state, False if not.
95    """
96    callback_received = uwb_test_utils.verify_uwb_state_callback(
97        dut, event_str, handler
98    )
99    if prev_uwb_state == expected_uwb_state:
100      return uwb_test_utils.get_uwb_state(dut) == expected_uwb_state
101    return callback_received
102
103  ### Test Cases ###
104
105  def test_toggle_uwb_state(self):
106    """Disables and verifies UWB state."""
107    uwb_test_utils.set_uwb_state_and_verify(self.dut, False, self.handler)
108    uwb_test_utils.set_uwb_state_and_verify(self.dut, True, self.handler)
109
110  def test_enable_uwb_state(self):
111    """Enables and verifies UWB state."""
112    uwb_test_utils.set_uwb_state_and_verify(self.dut, True, self.handler)
113
114  def test_uwb_state_after_reboot_with_uwb_off(self):
115    """Sets UWB state to off and verifies it is persistent after reboot."""
116    self._test_uwb_state_after_reboot(self.dut, False, self.handler)
117
118  def test_uwb_state_after_reboot_with_uwb_on(self):
119    """Sets UWB state to on and verifies it is persistent after reboot."""
120    self._test_uwb_state_after_reboot(self.dut, True, self.handler)
121
122  def test_uwb_state_with_airplane_mode_toggle(self):
123    """Verifies UWB is disabled with airplane mode on."""
124
125    # Enable APM. Verify UWB is disabled.
126    uwb_state = uwb_test_utils.get_uwb_state(self.dut)
127    uwb_test_utils.set_airplane_mode(self.dut, True)
128    asserts.assert_true(
129        self._verify_uwb_state_with_airplane_mode(
130            self.dut, uwb_state, False, "Disabled", self.handler
131        ),
132        "UWB is not disabled with airplane mode On.",
133    )
134
135    # Enable UWB with APM ON. Verify UWB is still disabled.
136    uwb_state = uwb_test_utils.get_uwb_state(self.dut)
137    self.dut.uwb.setUwbEnabled(True)
138    asserts.assert_true(
139        self._verify_uwb_state_with_airplane_mode(
140            self.dut, uwb_state, False, "Disabled", self.handler
141        ),
142        "Enabling UWB with airplane mode On should not work.",
143    )
144
145    # Disable APM. Verify UWB is enabled.
146    uwb_state = uwb_test_utils.get_uwb_state(self.dut)
147    uwb_test_utils.set_airplane_mode(self.dut, False)
148    asserts.assert_true(
149        self._verify_uwb_state_with_airplane_mode(
150            self.dut, uwb_state, True, "Inactive", self.handler
151        ),
152        "UWB is not enabled with airplane mode Off.",
153    )
154
155  def test_uwb_state_off_with_airplane_mode_toggle(self):
156    """Verifies UWB disabled state is persistent with airplane mode toggle."""
157
158    # disable UWB
159    uwb_test_utils.set_uwb_state_and_verify(self.dut, False, self.handler)
160
161    # enable airplane mode and verify UWB is disabled.
162    uwb_state = uwb_test_utils.get_uwb_state(self.dut)
163    uwb_test_utils.set_airplane_mode(self.dut, True)
164    asserts.assert_true(
165        self._verify_uwb_state_with_airplane_mode(
166            self.dut, uwb_state, False, "Disabled", self.handler
167        ),
168        "UWB is not disabled with airplane mode On.",
169    )
170
171    # disable airplane mode and verify UWB is disabled.
172    uwb_state = uwb_test_utils.get_uwb_state(self.dut)
173    uwb_test_utils.set_airplane_mode(self.dut, False)
174    asserts.assert_true(
175        self._verify_uwb_state_with_airplane_mode(
176            self.dut, uwb_state, False, "Disabled", self.handler
177        ),
178        "UWB state: On, Expected state: Off",
179    )
180
181    # enable UWB
182    uwb_test_utils.set_uwb_state_and_verify(self.dut, True)
183
184if __name__ == "__main__":
185  if "--" in sys.argv:
186    index = sys.argv.index("--")
187    sys.argv = sys.argv[:1] + sys.argv[index + 1:]
188  suite_runner.run_suite([UwbManagerTest])
189