commit fafcf1f66cd762d1d0b13da0190736254b41d955
author Roshan Pius <rpius@google.com> Wed Dec 15 11:08:10 2021 -0800
committer Roshan Pius <rpius@google.com> Wed Dec 15 11:11:59 2021 -0800
tree 89bac1a2224a3f4477bb641b24920085e45ec279
parent 543ccab19da31fe8f9eb0e2dadf503757c06e724

UwbManagerTest: Add a helper method for uwb toggle change

Bug: 209028991
Test: atest UwbManagerTest
Change-Id: Ia4f4ecba3fe6070d6214c849bb8d4ded355d920a

tests/uwb/src/android/uwb/cts/UwbManagerTest.java

diff --git a/tests/uwb/src/android/uwb/cts/UwbManagerTest.java b/tests/uwb/src/android/uwb/cts/UwbManagerTest.java
index 8447ece..489da9d 100644
--- a/tests/uwb/src/android/uwb/cts/UwbManagerTest.java
+++ b/tests/uwb/src/android/uwb/cts/UwbManagerTest.java
@@ -74,7 +74,7 @@
     private String mDefaultChipId;
 
     @Before
-    public void setup() {
+    public void setup() throws Exception {
         mUwbManager = mContext.getSystemService(UwbManager.class);
         assumeTrue(UwbTestUtils.isUwbSupported(mContext));
         assertThat(mUwbManager).isNotNull();
@@ -82,12 +82,30 @@
         // Ensure UWB is toggled on.
         ShellIdentityUtils.invokeWithShellPermissions(() -> {
             if (!mUwbManager.isUwbEnabled()) {
-                mUwbManager.setUwbEnabled(true);
+                try {
+                    setUwbEnabledAndWaitForCompletion(true);
+                } catch (Exception e) {
+                    fail("Exception while processing UWB toggle " + e);
+                }
             }
             mDefaultChipId = mUwbManager.getDefaultChipId();
         });
     }
 
+    // Should be invoked with shell permissions.
+    private void setUwbEnabledAndWaitForCompletion(boolean enabled) throws Exception {
+        CountDownLatch countDownLatch = new CountDownLatch(1);
+        int adapterState = enabled ? STATE_ENABLED_INACTIVE : STATE_DISABLED;
+        AdapterStateCallback adapterStateCallback =
+                new AdapterStateCallback(countDownLatch, adapterState);
+        mUwbManager.registerAdapterStateCallback(
+                Executors.newSingleThreadExecutor(), adapterStateCallback);
+        mUwbManager.setUwbEnabled(enabled);
+        assertThat(countDownLatch.await(2, TimeUnit.SECONDS)).isTrue();
+        assertThat(mUwbManager.isUwbEnabled()).isEqualTo(enabled);
+        assertThat(adapterStateCallback.state).isEqualTo(adapterState);
+    }
+
     @Test
     public void testGetSpecificationInfo() {
         UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
@@ -521,22 +539,13 @@
     @Test
     public void testUwbStateToggle() throws Exception {
         UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
-        CountDownLatch countDownLatch = new CountDownLatch(1);
-        AdapterStateCallback adapterStateCallback =
-                new AdapterStateCallback(countDownLatch, STATE_DISABLED);
         try {
             // Needs UWB_PRIVILEGED permission which is held by shell.
             uiAutomation.adoptShellPermissionIdentity();
-            mUwbManager.registerAdapterStateCallback(
-                    Executors.newSingleThreadExecutor(), adapterStateCallback);
             assertThat(mUwbManager.isUwbEnabled()).isTrue();
             assertThat(mUwbManager.getAdapterState()).isEqualTo(STATE_ENABLED_INACTIVE);
             // Toggle the state
-            mUwbManager.setUwbEnabled(false);
-            // Wait for the on start failed callback.
-            assertThat(countDownLatch.await(1, TimeUnit.SECONDS)).isTrue();
-            assertThat(mUwbManager.isUwbEnabled()).isFalse();
-            assertThat(adapterStateCallback.state).isEqualTo(STATE_DISABLED);
+            setUwbEnabledAndWaitForCompletion(false);
             assertThat(mUwbManager.getAdapterState()).isEqualTo(STATE_DISABLED);
         } finally {
             uiAutomation.dropShellPermissionIdentity();