Make getAdapterState() consistent with AdapterState notification.

When the country code is not valid, or, it's valid but setting it
results in an error, the AdapterState notification is sent with
STATE_DISABLED.

This CL updates the UwbServiceCore.getAdapterState() API to return the same state as in the AdapterState notification (STATE_DISABLED in the above cases).

Test: atest ServiceUwbTests
Bug: 276383595
Change-Id: I18f64a3dfab54fbf466765b62b548b32276c1372
diff --git a/service/java/com/android/server/uwb/UwbServiceCore.java b/service/java/com/android/server/uwb/UwbServiceCore.java
index 45b6893..b841822 100644
--- a/service/java/com/android/server/uwb/UwbServiceCore.java
+++ b/service/java/com/android/server/uwb/UwbServiceCore.java
@@ -188,7 +188,7 @@
 
     private boolean isUwbEnabled() {
         synchronized (UwbServiceCore.this) {
-            return getAdapterState() != AdapterStateCallback.STATE_DISABLED;
+            return getInternalAdapterState() != AdapterStateCallback.STATE_DISABLED;
         }
     }
 
@@ -242,11 +242,7 @@
             return;
         }
         updateDeviceState(deviceState, chipId);
-        // TODO(b/244443764): We should use getAdapterState() here, as for multi-chip case
-        // the configured state above can be different from the computed adapter state
-        // (after the call to updateState()).
         mUwbTask.computeAndNotifyAdapterStateChange(
-                getAdapterStateFromDeviceState(deviceState),
                 getReasonFromDeviceState(deviceState),
                 mUwbCountryCode.getCountryCode(),
                 Optional.empty());
@@ -254,7 +250,7 @@
 
     void updateDeviceState(int deviceState, String chipId) {
         Log.i(TAG, "updateState(): deviceState = " + getDeviceStateString(deviceState)
-                + ", current adapter state = " + getAdapterState());
+                + ", current internal adapter state = " + getInternalAdapterState());
 
         oemExtensionDeviceStatusUpdate(deviceState, chipId);
         updateState(
@@ -346,7 +342,7 @@
             // Notify the current UWB adapter state. For example, if UWB was earlier enabled and at
             // that time the country code was not valid, will now notify STATE_ENABLED_INACTIVE.
             mUwbTask.computeAndNotifyAdapterStateChange(
-                    getAdapterState(), mLastStateChangedReason, countryCode, Optional.empty());
+                    mLastStateChangedReason, countryCode, Optional.empty());
         }
     }
 
@@ -572,7 +568,26 @@
         mSessionManager.sendData(sessionHandle, remoteDeviceAddress, params, data);
     }
 
+    /**
+     * Get the UWB Adapter State.
+     */
     public /* @UwbManager.AdapterStateCallback.State */ int getAdapterState() {
+        return computeAdapterState(mUwbCountryCode.getCountryCode(), Optional.empty());
+    }
+
+    private int computeAdapterState(String countryCode, Optional<Integer> setCountryCodeStatus) {
+        // When either the country code is not valid or setting it in UWBS failed with an error,
+        // notify the UWB stack state as DISABLED (even though internally the UWB device state
+        // may be stored as READY), so that applications wait for starting a ranging session.
+        if (!mUwbCountryCode.isValid(countryCode)
+                || (setCountryCodeStatus.isPresent()
+                && setCountryCodeStatus.get() != STATUS_CODE_OK)) {
+            return AdapterStateCallback.STATE_DISABLED;
+        }
+        return getInternalAdapterState();
+    }
+
+    private /* @UwbManager.AdapterStateCallback.State */ int getInternalAdapterState() {
         synchronized (UwbServiceCore.this) {
             if (mChipIdToStateMap.isEmpty()) {
                 return AdapterStateCallback.STATE_DISABLED;
@@ -809,7 +824,6 @@
                         String countryCode = setCountryCodeResult.second;
                         Log.i(TAG, "Current country code = " + countryCode);
                         computeAndNotifyAdapterStateChange(
-                                getAdapterStateFromDeviceState(UwbUciConstants.DEVICE_STATE_READY),
                                 getReasonFromDeviceState(UwbUciConstants.DEVICE_STATE_READY),
                                 countryCode,
                                 setCountryCodeStatus);
@@ -876,17 +890,8 @@
             }
         }
 
-        private void computeAndNotifyAdapterStateChange(int adapterState, int reason,
+        private void computeAndNotifyAdapterStateChange(int reason,
                 String countryCode, Optional<Integer> setCountryCodeStatus) {
-            // When either the country code is not valid or setting it in UWBS failed with an error,
-            // notify the UWB stack state as DISABLED (even though internally the UWB device state
-            // may be stored as READY), so that applications wait for starting a ranging session.
-            if (!mUwbCountryCode.isValid(countryCode)
-                    || (setCountryCodeStatus.isPresent()
-                        && setCountryCodeStatus.get() != STATUS_CODE_OK)) {
-                adapterState = AdapterStateCallback.STATE_DISABLED;
-            }
-
             // When either the country code is not valid or setting it in UWBS failed with the error
             // STATUS_CODE_ANDROID_REGULATION_UWB_OFF, notify with the reason SYSTEM_REGULATION.
             if (!mUwbCountryCode.isValid(countryCode)
@@ -896,7 +901,7 @@
                 reason = StateChangeReason.SYSTEM_REGULATION;
             }
 
-            notifyAdapterState(adapterState, reason);
+            notifyAdapterState(computeAdapterState(countryCode, setCountryCodeStatus), reason);
         }
 
         public class WatchDogThread extends Thread {
@@ -956,6 +961,7 @@
                     + " for chip id = " + chipId);
         }
         pw.println("mLastStateChangedReason = " + mLastStateChangedReason);
+        pw.println("mLastAdapterStateNotification = " + mLastAdapterStateNotification);
         pw.println("---- Dump of UwbServiceCore ----");
     }
 
@@ -996,7 +1002,7 @@
 
             return new UwbActivityEnergyInfo.Builder()
                     .setTimeSinceBootMillis(SystemClock.elapsedRealtime())
-                    .setStackState(getAdapterState())
+                    .setStackState(getInternalAdapterState())
                     .setControllerTxDurationMillis(stats.getTxTimeMs())
                     .setControllerRxDurationMillis(stats.getRxTimeMs())
                     .setControllerIdleDurationMillis(stats.getIdleTimeMs())
diff --git a/service/tests/src/com/android/server/uwb/UwbServiceCoreTest.java b/service/tests/src/com/android/server/uwb/UwbServiceCoreTest.java
index 89c5590..bb2d878 100644
--- a/service/tests/src/com/android/server/uwb/UwbServiceCoreTest.java
+++ b/service/tests/src/com/android/server/uwb/UwbServiceCoreTest.java
@@ -285,9 +285,9 @@
                 genericSpecificationParams);
     }
 
-    private void enableUwb() throws Exception {
+    private void enableUwb(String countryCode) throws Exception {
         when(mNativeUwbManager.doInitialize()).thenReturn(true);
-        when(mUwbCountryCode.getCountryCode()).thenReturn(null);
+        when(mUwbCountryCode.getCountryCode()).thenReturn(countryCode);
         when(mUwbCountryCode.setCountryCode(anyBoolean())).thenReturn(
                 Pair.create(STATUS_CODE_FAILED, null));
 
@@ -349,13 +349,14 @@
         clearInvocations(cb);
 
         // Enable (with country code initially unknown, like at boot time).
-        enableUwb();
+        enableUwb(null);
 
         verify(mNativeUwbManager).doInitialize();
         verify(mUwbCountryCode).setCountryCode(true);
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_REGULATION);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
 
         clearInvocations(cb);
 
@@ -363,14 +364,19 @@
         // country code. We don't expect any more AdapterState notifications as the Adapter State
         // is still considered to be the same (STATE_DISABLED).
         mUwbServiceCore.onCountryCodeChanged("00");
+
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
 
         // Valid country code changed notification is received after some time (before the timeout).
         // The message queue immediately has a message to process, which results in a call to the
         // adapter state callback.
+        when(mUwbCountryCode.getCountryCode()).thenReturn(VALID_COUNTRY_CODE);
         mUwbServiceCore.onCountryCodeChanged("US");
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
     }
 
     // Unit test for scenario when setting the country code (during UWB Enable) fails with a UWB
@@ -389,7 +395,7 @@
         // Enable (with country code that results in the Vendor-specific UWB_REGULATION error).
         when(mNativeUwbManager.doInitialize()).thenReturn(true);
 
-        when(mUwbCountryCode.getCountryCode()).thenReturn("JP");
+        when(mUwbCountryCode.getCountryCode()).thenReturn(null);
         when(mUwbCountryCode.setCountryCode(anyBoolean()))
                 .thenReturn(Pair.create(STATUS_CODE_ANDROID_REGULATION_UWB_OFF, "JP"));
 
@@ -400,6 +406,7 @@
         verify(mUwbCountryCode).setCountryCode(true);
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_REGULATION);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     // Unit test for scenario when setting the country code (during UWB Enable) fails with a generic
@@ -415,9 +422,10 @@
 
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
-        // Enable (with a valid country code), but the firmware returns some error.
+        // Enable (with a valid country code), but the firmware returns some error. Hence,
+        // UwbCountryCode.getCountryCode() will return null.
         when(mNativeUwbManager.doInitialize()).thenReturn(true);
-        when(mUwbCountryCode.getCountryCode()).thenReturn(VALID_COUNTRY_CODE);
+        when(mUwbCountryCode.getCountryCode()).thenReturn(null);
         when(mUwbCountryCode.setCountryCode(anyBoolean())).thenReturn(Pair.create(
                 STATUS_CODE_FAILED, VALID_COUNTRY_CODE));
 
@@ -428,6 +436,7 @@
         verify(mUwbCountryCode).setCountryCode(true);
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     @Test
@@ -441,16 +450,21 @@
         // Enabled UWB, we expect an Adapter State notification with State ENABLED_INACTIVE as
         // there is a valid country code.
         enableUwbWithCountryCodeChangedCallback();
+
         verify(mNativeUwbManager).doInitialize();
         verify(mUwbCountryCode).setCountryCode(true);
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
 
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
         // Enable again. should be ignored.
-        enableUwb();
+        enableUwb(VALID_COUNTRY_CODE);
         verifyNoMoreInteractions(mNativeUwbManager, mUwbCountryCode, cb);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
     }
 
     // Test the UWB stack enable when the NativeUwbManager.doInitialize() is delayed such that the
@@ -494,6 +508,8 @@
         verify(mUwbWakeLock, times(1)).release();
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
     }
 
     @Test
@@ -512,6 +528,9 @@
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
+
         clearInvocations(cb);
 
         // Disable UWB.
@@ -520,6 +539,7 @@
         verify(mNativeUwbManager).doDeinitialize();
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     // Test the UWB stack disable when the NativeUwbManager.doDeinitialize() is delayed such that
@@ -534,8 +554,11 @@
 
         // Enable first
         enableUwbWithCountryCodeChangedCallback();
+
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
 
         clearInvocations(mUwbWakeLock, cb);
 
@@ -562,6 +585,7 @@
         verify(mUwbWakeLock, times(1)).release();
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_POLICY);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     @Test
@@ -575,23 +599,32 @@
 
         // Enable first
         enableUwbWithCountryCodeChangedCallback();
+
         verify(mNativeUwbManager).doInitialize();
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
+
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
         // Disable UWB.
         disableUwb();
+
         verify(mNativeUwbManager).doDeinitialize();
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(mUwbCountryCode);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
+
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
         // Disable again. should be ignored.
         disableUwb();
+
         verifyNoMoreInteractions(mNativeUwbManager, mUwbCountryCode, cb);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     @Test
@@ -607,22 +640,26 @@
         clearInvocations(cb);
 
         // Enable first (with country code initially unknown, like at boot time).
-        enableUwb();
+        enableUwb(null);
 
         verify(mNativeUwbManager).doInitialize();
         verify(mUwbCountryCode).setCountryCode(true);
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_REGULATION);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
 
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
         // Later on, the onCountryCodeChanged() callback is invoked, with a valid country code.
+        when(mUwbCountryCode.getCountryCode()).thenReturn(VALID_COUNTRY_CODE);
         mUwbServiceCore.onCountryCodeChanged(VALID_COUNTRY_CODE);
 
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
 
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
@@ -633,6 +670,7 @@
         verify(cb).onAdapterStateChanged(UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
 
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
@@ -646,6 +684,8 @@
                 UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
     }
 
     @Test
@@ -690,6 +730,7 @@
                 UwbManager.AdapterStateCallback.STATE_ENABLED_INACTIVE,
                 StateChangeReason.SYSTEM_POLICY);
         verifyNoMoreInteractions(cb, mNativeUwbManager);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     @Test
@@ -940,7 +981,7 @@
         assertThat(mUwbServiceCore.getAdapterState())
                 .isEqualTo(AdapterState.STATE_ENABLED_INACTIVE);
 
-        // If one chip is active, then getAdapterState should return STATE_ENABLED_ACTIVE.
+        // If one chip is active, then getAdapterState() should return STATE_ENABLED_ACTIVE.
         mUwbServiceCore.onDeviceStatusNotificationReceived(UwbUciConstants.DEVICE_STATE_ACTIVE,
                 TEST_CHIP_ONE_CHIP_ID);
         mTestLooper.dispatchAll();
@@ -965,7 +1006,8 @@
                 mUwbCountryCode, mUwbSessionManager, mUwbConfigurationManager,
                 mUwbInjector, mTestLooper.getLooper());
 
-        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
+        assertThat(mUwbServiceCore.getAdapterState())
+                .isEqualTo(AdapterState.STATE_DISABLED);
     }
 
 
@@ -1192,11 +1234,14 @@
 
         // UWB Adapter State will internally be Enabled, but we expect an AdapterState notification
         // with State DISABLED to be sent (as the country code is not valid).
-        enableUwb();
+        enableUwb(null);
+
         verify(cb).onAdapterStateChanged(
                 UwbManager.AdapterStateCallback.STATE_DISABLED,
                 StateChangeReason.SYSTEM_REGULATION);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
         verifyNoMoreInteractions(cb);
+
         clearInvocations(mNativeUwbManager, mUwbCountryCode, cb);
 
         // UWB Enable is expected to result in setting the country code; but that may fail and the
@@ -1216,6 +1261,7 @@
         verify(mNativeUwbManager).doDeinitialize();
         verify(mNativeUwbManager).doInitialize();
         verifyNoMoreInteractions(cb);
+        assertThat(mUwbServiceCore.getAdapterState()).isEqualTo(AdapterState.STATE_DISABLED);
     }
 
     @Test