2 files changed, 4 insertions, 9 deletions
diff --git a/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
index 8d5e01c5b782..9050dad0134c 100644
--- a/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
+++ b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
@@ -29,13 +29,14 @@ import android.os.Bundle
 import android.os.Handler
 import android.os.UserHandle
 import android.util.Log
-import android.view.InputDevice
 import androidx.core.app.NotificationCompat
 import androidx.core.app.NotificationManagerCompat
 import com.android.internal.annotations.VisibleForTesting
 import com.android.systemui.R
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.shared.hardware.hasInputDevice
+import com.android.systemui.shared.hardware.isAnyStylusSource
 import com.android.systemui.util.NotificationChannels
 import java.text.NumberFormat
 import javax.inject.Inject
@@ -150,10 +151,7 @@ constructor(
     }
 
     private fun hasConnectedBluetoothStylus(): Boolean {
-        // TODO(b/257936830): get bt address once input api available
-        return inputManager.inputDeviceIds.any { deviceId ->
-            inputManager.getInputDevice(deviceId).supportsSource(InputDevice.SOURCE_STYLUS)
-        }
+        return inputManager.hasInputDevice { it.isAnyStylusSource && it.bluetoothAddress != null }
     }
 
     private fun getPendingBroadcast(action: String): PendingIntent? {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt
index 1e81dc761b18..e1668e8bfa37 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt
@@ -36,7 +36,6 @@ import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.assertEquals
 import org.junit.Before
-import org.junit.Ignore
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
@@ -79,7 +78,7 @@ class StylusUsiPowerUiTest : SysuiTestCase() {
         whenever(inputManager.inputDeviceIds).thenReturn(intArrayOf())
         whenever(inputManager.getInputDevice(0)).thenReturn(btStylusDevice)
         whenever(btStylusDevice.supportsSource(InputDevice.SOURCE_STYLUS)).thenReturn(true)
-        // whenever(btStylusDevice.bluetoothAddress).thenReturn("SO:ME:AD:DR:ES")
+        whenever(btStylusDevice.bluetoothAddress).thenReturn("SO:ME:AD:DR:ES")
 
         stylusUsiPowerUi = StylusUsiPowerUI(contextSpy, notificationManager, inputManager, handler)
         broadcastReceiver = stylusUsiPowerUi.receiver
@@ -179,7 +178,6 @@ class StylusUsiPowerUiTest : SysuiTestCase() {
     }
 
     @Test
-    @Ignore("TODO(b/257936830): get bt address once input api available")
     fun refresh_hasConnectedBluetoothStylus_cancelsNotification() {
         whenever(inputManager.inputDeviceIds).thenReturn(intArrayOf(0))
 
@@ -189,7 +187,6 @@ class StylusUsiPowerUiTest : SysuiTestCase() {
     }
 
     @Test
-    @Ignore("TODO(b/257936830): get bt address once input api available")
     fun refresh_hasConnectedBluetoothStylus_existingNotification_cancelsNotification() {
         stylusUsiPowerUi.updateBatteryState(0, FixedCapacityBatteryState(0.1f))
         whenever(inputManager.inputDeviceIds).thenReturn(intArrayOf(0))