Merge "Low battery notifications for USI styluses." into tm-qpr-dev

7 files changed, 646 insertions, 0 deletions

diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 9d5a8c4d3cc7..ca2cf1a56b41 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -2762,4 +2762,7 @@
     <string name="rear_display_accessibility_folded_animation">Foldable device being unfolded</string>
     <!-- Text for education page content description for unfolded animation. [CHAR_LIMIT=NONE] -->
     <string name="rear_display_accessibility_unfolded_animation">Foldable device being flipped around</string>
+
+    <!-- Title for notification of low stylus battery. [CHAR_LIMIT=NONE] -->
+    <string name="stylus_battery_low">Stylus battery low</string>
 </resources>
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java b/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java
index 7864f1901e57..ccb9557d757b 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java
@@ -93,6 +93,8 @@ import android.view.accessibility.AccessibilityManager;
 import android.view.accessibility.CaptioningManager;
 import android.view.inputmethod.InputMethodManager;
 
+import androidx.core.app.NotificationManagerCompat;
+
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.appwidget.IAppWidgetService;
 import com.android.internal.jank.InteractionJankMonitor;
@@ -389,6 +391,12 @@ public class FrameworkServicesModule {
         return context.getSystemService(NotificationManager.class);
     }
 
+    @Provides
+    @Singleton
+    static NotificationManagerCompat provideNotificationManagerCompat(Context context) {
+        return NotificationManagerCompat.from(context);
+    }
+
     /** */
     @Provides
     @Singleton
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt b/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
index 0fbe0acb5642..a4252bbe1d68 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
@@ -42,6 +42,7 @@ import com.android.systemui.settings.dagger.MultiUserUtilsModule
 import com.android.systemui.shortcut.ShortcutKeyDispatcher
 import com.android.systemui.statusbar.notification.InstantAppNotifier
 import com.android.systemui.statusbar.phone.KeyguardLiftController
+import com.android.systemui.stylus.StylusUsiPowerStartable
 import com.android.systemui.temporarydisplay.chipbar.ChipbarCoordinator
 import com.android.systemui.theme.ThemeOverlayController
 import com.android.systemui.toast.ToastUI
@@ -251,4 +252,10 @@ abstract class SystemUICoreStartableModule {
     @IntoMap
     @ClassKey(RearDisplayDialogController::class)
     abstract fun bindRearDisplayDialogController(sysui: RearDisplayDialogController): CoreStartable
+
+    /** Inject into StylusUsiPowerStartable) */
+    @Binds
+    @IntoMap
+    @ClassKey(StylusUsiPowerStartable::class)
+    abstract fun bindStylusUsiPowerStartable(sysui: StylusUsiPowerStartable): CoreStartable
 }
diff --git a/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerStartable.kt b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerStartable.kt
new file mode 100644
index 000000000000..11233dda165c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerStartable.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.stylus
+
+import android.hardware.BatteryState
+import android.hardware.input.InputManager
+import android.util.Log
+import android.view.InputDevice
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
+import java.util.concurrent.Executor
+import javax.inject.Inject
+
+/**
+ * A [CoreStartable] that listens to USI stylus battery events, to manage the [StylusUsiPowerUI]
+ * notification controller.
+ */
+@SysUISingleton
+class StylusUsiPowerStartable
+@Inject
+constructor(
+    private val stylusManager: StylusManager,
+    private val inputManager: InputManager,
+    private val stylusUsiPowerUi: StylusUsiPowerUI,
+    private val featureFlags: FeatureFlags,
+    @Background private val executor: Executor,
+) : CoreStartable, StylusManager.StylusCallback, InputManager.InputDeviceBatteryListener {
+
+    override fun onStylusAdded(deviceId: Int) {
+        val device = inputManager.getInputDevice(deviceId) ?: return
+
+        if (!device.isExternal) {
+            registerBatteryListener(deviceId)
+        }
+    }
+
+    override fun onStylusBluetoothConnected(deviceId: Int, btAddress: String) {
+        stylusUsiPowerUi.refresh()
+    }
+
+    override fun onStylusBluetoothDisconnected(deviceId: Int, btAddress: String) {
+        stylusUsiPowerUi.refresh()
+    }
+
+    override fun onStylusRemoved(deviceId: Int) {
+        val device = inputManager.getInputDevice(deviceId) ?: return
+
+        if (!device.isExternal) {
+            unregisterBatteryListener(deviceId)
+        }
+    }
+
+    override fun onBatteryStateChanged(
+        deviceId: Int,
+        eventTimeMillis: Long,
+        batteryState: BatteryState
+    ) {
+        if (batteryState.isPresent) {
+            stylusUsiPowerUi.updateBatteryState(batteryState)
+        }
+    }
+
+    private fun registerBatteryListener(deviceId: Int) {
+        try {
+            inputManager.addInputDeviceBatteryListener(deviceId, executor, this)
+        } catch (e: SecurityException) {
+            Log.e(TAG, "$e: Failed to register battery listener for $deviceId.")
+        }
+    }
+
+    private fun unregisterBatteryListener(deviceId: Int) {
+        try {
+            inputManager.removeInputDeviceBatteryListener(deviceId, this)
+        } catch (e: SecurityException) {
+            Log.e(TAG, "$e: Failed to unregister battery listener for $deviceId.")
+        }
+    }
+
+    override fun start() {
+        if (!featureFlags.isEnabled(Flags.ENABLE_USI_BATTERY_NOTIFICATIONS)) return
+        addBatteryListenerForInternalStyluses()
+
+        stylusManager.registerCallback(this)
+        stylusManager.startListener()
+    }
+
+    private fun addBatteryListenerForInternalStyluses() {
+        // For most devices, an active stylus is represented by an internal InputDevice.
+        // This InputDevice will be present in InputManager before CoreStartables run,
+        // and will not be removed. In many cases, it reports the battery level of the stylus.
+        inputManager.inputDeviceIds
+            .asSequence()
+            .mapNotNull { inputManager.getInputDevice(it) }
+            .filter { it.supportsSource(InputDevice.SOURCE_STYLUS) }
+            .forEach { onStylusAdded(it.id) }
+    }
+
+    companion object {
+        private val TAG = StylusUsiPowerStartable::class.simpleName.orEmpty()
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
new file mode 100644
index 000000000000..70a5b366263e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
@@ -0,0 +1,185 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.stylus
+
+import android.Manifest
+import android.app.PendingIntent
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.hardware.BatteryState
+import android.hardware.input.InputManager
+import android.os.Handler
+import android.os.UserHandle
+import android.view.InputDevice
+import androidx.core.app.NotificationCompat
+import androidx.core.app.NotificationManagerCompat
+import com.android.systemui.R
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.util.NotificationChannels
+import java.text.NumberFormat
+import javax.inject.Inject
+
+/**
+ * UI controller for the notification that shows when a USI stylus battery is low. The
+ * [StylusUsiPowerStartable], which listens to battery events, uses this controller.
+ */
+@SysUISingleton
+class StylusUsiPowerUI
+@Inject
+constructor(
+    private val context: Context,
+    private val notificationManager: NotificationManagerCompat,
+    private val inputManager: InputManager,
+    @Background private val handler: Handler,
+) {
+
+    // These values must only be accessed on the handler.
+    private var batteryCapacity = 1.0f
+    private var suppressed = false
+
+    fun init() {
+        val filter =
+            IntentFilter().also {
+                it.addAction(ACTION_DISMISSED_LOW_BATTERY)
+                it.addAction(ACTION_CLICKED_LOW_BATTERY)
+            }
+
+        context.registerReceiverAsUser(
+            receiver,
+            UserHandle.ALL,
+            filter,
+            Manifest.permission.DEVICE_POWER,
+            handler,
+            Context.RECEIVER_NOT_EXPORTED,
+        )
+    }
+
+    fun refresh() {
+        handler.post refreshNotification@{
+            if (!suppressed && !hasConnectedBluetoothStylus() && isBatteryBelowThreshold()) {
+                showOrUpdateNotification()
+                return@refreshNotification
+            }
+
+            if (!isBatteryBelowThreshold()) {
+                // Reset suppression when stylus battery is recharged, so that the next time
+                // it reaches a low battery, the notification will show again.
+                suppressed = false
+            }
+            hideNotification()
+        }
+    }
+
+    fun updateBatteryState(batteryState: BatteryState) {
+        handler.post updateBattery@{
+            if (batteryState.capacity == batteryCapacity) return@updateBattery
+
+            batteryCapacity = batteryState.capacity
+            refresh()
+        }
+    }
+
+    /**
+     * Suppression happens when the notification is dismissed by the user. This is to prevent
+     * further battery events with capacities below the threshold from reopening the suppressed
+     * notification.
+     *
+     * Suppression can only be removed when the battery has been recharged - thus restarting the
+     * notification cycle (i.e. next low battery event, notification should show).
+     */
+    fun updateSuppression(suppress: Boolean) {
+        handler.post updateSuppressed@{
+            if (suppressed == suppress) return@updateSuppressed
+
+            suppressed = suppress
+            refresh()
+        }
+    }
+
+    private fun hideNotification() {
+        notificationManager.cancel(USI_NOTIFICATION_ID)
+    }
+
+    private fun showOrUpdateNotification() {
+        val notification =
+            NotificationCompat.Builder(context, NotificationChannels.BATTERY)
+                .setSmallIcon(R.drawable.ic_power_low)
+                .setDeleteIntent(getPendingBroadcast(ACTION_DISMISSED_LOW_BATTERY))
+                .setContentIntent(getPendingBroadcast(ACTION_CLICKED_LOW_BATTERY))
+                .setContentTitle(context.getString(R.string.stylus_battery_low))
+                .setContentText(
+                    context.getString(
+                        R.string.battery_low_percent_format,
+                        NumberFormat.getPercentInstance().format(batteryCapacity)
+                    )
+                )
+                .setPriority(NotificationCompat.PRIORITY_DEFAULT)
+                .setLocalOnly(true)
+                .setAutoCancel(true)
+                .build()
+
+        notificationManager.notify(USI_NOTIFICATION_ID, notification)
+    }
+
+    private fun isBatteryBelowThreshold(): Boolean {
+        return batteryCapacity <= LOW_BATTERY_THRESHOLD
+    }
+
+    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)
+        }
+    }
+
+    private fun getPendingBroadcast(action: String): PendingIntent? {
+        return PendingIntent.getBroadcastAsUser(
+            context,
+            0,
+            Intent(action),
+            PendingIntent.FLAG_IMMUTABLE,
+            UserHandle.CURRENT
+        )
+    }
+
+    private val receiver: BroadcastReceiver =
+        object : BroadcastReceiver() {
+            override fun onReceive(context: Context, intent: Intent) {
+                when (intent.action) {
+                    ACTION_DISMISSED_LOW_BATTERY -> updateSuppression(true)
+                    ACTION_CLICKED_LOW_BATTERY -> {
+                        updateSuppression(true)
+                        // TODO(b/261584943): open USI device details page
+                    }
+                }
+            }
+        }
+
+    companion object {
+        // Low battery threshold matches CrOS, see:
+        // https://source.chromium.org/chromium/chromium/src/+/main:ash/system/power/peripheral_battery_notifier.cc;l=41
+        private const val LOW_BATTERY_THRESHOLD = 0.16f
+
+        private val USI_NOTIFICATION_ID = R.string.stylus_battery_low
+
+        private const val ACTION_DISMISSED_LOW_BATTERY = "StylusUsiPowerUI.dismiss"
+        private const val ACTION_CLICKED_LOW_BATTERY = "StylusUsiPowerUI.click"
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerStartableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerStartableTest.kt
new file mode 100644
index 000000000000..ff382a3ec19f
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerStartableTest.kt
@@ -0,0 +1,158 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.stylus
+
+import android.hardware.BatteryState
+import android.hardware.input.InputManager
+import android.testing.AndroidTestingRunner
+import android.view.InputDevice
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
+import com.android.systemui.util.mockito.whenever
+import java.util.concurrent.Executor
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.inOrder
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyNoMoreInteractions
+import org.mockito.MockitoAnnotations
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class StylusUsiPowerStartableTest : SysuiTestCase() {
+    @Mock lateinit var inputManager: InputManager
+    @Mock lateinit var stylusManager: StylusManager
+    @Mock lateinit var stylusDevice: InputDevice
+    @Mock lateinit var externalDevice: InputDevice
+    @Mock lateinit var featureFlags: FeatureFlags
+    @Mock lateinit var stylusUsiPowerUi: StylusUsiPowerUI
+
+    lateinit var startable: StylusUsiPowerStartable
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        startable =
+            StylusUsiPowerStartable(
+                stylusManager,
+                inputManager,
+                stylusUsiPowerUi,
+                featureFlags,
+                DIRECT_EXECUTOR,
+            )
+
+        whenever(featureFlags.isEnabled(Flags.ENABLE_USI_BATTERY_NOTIFICATIONS)).thenReturn(true)
+
+        whenever(inputManager.getInputDevice(EXTERNAL_DEVICE_ID)).thenReturn(externalDevice)
+        whenever(inputManager.getInputDevice(STYLUS_DEVICE_ID)).thenReturn(stylusDevice)
+        whenever(inputManager.inputDeviceIds)
+            .thenReturn(intArrayOf(EXTERNAL_DEVICE_ID, STYLUS_DEVICE_ID))
+
+        whenever(stylusDevice.supportsSource(InputDevice.SOURCE_STYLUS)).thenReturn(true)
+        whenever(stylusDevice.isExternal).thenReturn(false)
+        whenever(stylusDevice.id).thenReturn(STYLUS_DEVICE_ID)
+        whenever(externalDevice.supportsSource(InputDevice.SOURCE_STYLUS)).thenReturn(true)
+        whenever(externalDevice.isExternal).thenReturn(true)
+        whenever(externalDevice.id).thenReturn(EXTERNAL_DEVICE_ID)
+    }
+
+    @Test
+    fun start_addsBatteryListenerForInternalStylus() {
+        startable.start()
+
+        verify(inputManager, times(1))
+            .addInputDeviceBatteryListener(STYLUS_DEVICE_ID, DIRECT_EXECUTOR, startable)
+    }
+
+    @Test
+    fun onStylusAdded_internalStylus_addsBatteryListener() {
+        startable.onStylusAdded(STYLUS_DEVICE_ID)
+
+        verify(inputManager, times(1))
+            .addInputDeviceBatteryListener(STYLUS_DEVICE_ID, DIRECT_EXECUTOR, startable)
+    }
+
+    @Test
+    fun onStylusAdded_externalStylus_doesNotAddBatteryListener() {
+        startable.onStylusAdded(EXTERNAL_DEVICE_ID)
+
+        verify(inputManager, never())
+            .addInputDeviceBatteryListener(EXTERNAL_DEVICE_ID, DIRECT_EXECUTOR, startable)
+    }
+
+    @Test
+    fun onStylusRemoved_registeredStylus_removesBatteryListener() {
+        startable.onStylusAdded(STYLUS_DEVICE_ID)
+        startable.onStylusRemoved(STYLUS_DEVICE_ID)
+
+        inOrder(inputManager).let {
+            it.verify(inputManager, times(1))
+                .addInputDeviceBatteryListener(STYLUS_DEVICE_ID, DIRECT_EXECUTOR, startable)
+            it.verify(inputManager, times(1))
+                .removeInputDeviceBatteryListener(STYLUS_DEVICE_ID, startable)
+        }
+    }
+
+    @Test
+    fun onStylusBluetoothConnected_refreshesNotification() {
+        startable.onStylusBluetoothConnected(STYLUS_DEVICE_ID, "ANY")
+
+        verify(stylusUsiPowerUi, times(1)).refresh()
+    }
+
+    @Test
+    fun onStylusBluetoothDisconnected_refreshesNotification() {
+        startable.onStylusBluetoothDisconnected(STYLUS_DEVICE_ID, "ANY")
+
+        verify(stylusUsiPowerUi, times(1)).refresh()
+    }
+
+    @Test
+    fun onBatteryStateChanged_batteryPresent_refreshesNotification() {
+        val batteryState = mock(BatteryState::class.java)
+        whenever(batteryState.isPresent).thenReturn(true)
+
+        startable.onBatteryStateChanged(STYLUS_DEVICE_ID, 123, batteryState)
+
+        verify(stylusUsiPowerUi, times(1)).updateBatteryState(batteryState)
+    }
+
+    @Test
+    fun onBatteryStateChanged_batteryNotPresent_noop() {
+        val batteryState = mock(BatteryState::class.java)
+        whenever(batteryState.isPresent).thenReturn(false)
+
+        startable.onBatteryStateChanged(STYLUS_DEVICE_ID, 123, batteryState)
+
+        verifyNoMoreInteractions(stylusUsiPowerUi)
+    }
+
+    companion object {
+        private val DIRECT_EXECUTOR = Executor { r -> r.run() }
+
+        private const val EXTERNAL_DEVICE_ID = 0
+        private const val STYLUS_DEVICE_ID = 1
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt
new file mode 100644
index 000000000000..59875507341d
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt
@@ -0,0 +1,167 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.stylus
+
+import android.hardware.BatteryState
+import android.hardware.input.InputManager
+import android.os.Handler
+import android.testing.AndroidTestingRunner
+import android.view.InputDevice
+import androidx.core.app.NotificationManagerCompat
+import androidx.test.filters.SmallTest
+import com.android.systemui.R
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.whenever
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.inOrder
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyNoMoreInteractions
+import org.mockito.MockitoAnnotations
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class StylusUsiPowerUiTest : SysuiTestCase() {
+    @Mock lateinit var notificationManager: NotificationManagerCompat
+    @Mock lateinit var inputManager: InputManager
+    @Mock lateinit var handler: Handler
+    @Mock lateinit var btStylusDevice: InputDevice
+
+    private lateinit var stylusUsiPowerUi: StylusUsiPowerUI
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        whenever(handler.post(any())).thenAnswer {
+            (it.arguments[0] as Runnable).run()
+            true
+        }
+
+        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")
+
+        stylusUsiPowerUi = StylusUsiPowerUI(mContext, notificationManager, inputManager, handler)
+    }
+
+    @Test
+    fun updateBatteryState_capacityBelowThreshold_notifies() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+
+        verify(notificationManager, times(1)).notify(eq(R.string.stylus_battery_low), any())
+        verifyNoMoreInteractions(notificationManager)
+    }
+
+    @Test
+    fun updateBatteryState_capacityAboveThreshold_cancelsNotificattion() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.8f))
+
+        verify(notificationManager, times(1)).cancel(R.string.stylus_battery_low)
+        verifyNoMoreInteractions(notificationManager)
+    }
+
+    @Test
+    fun updateBatteryState_existingNotification_capacityAboveThreshold_cancelsNotification() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.8f))
+
+        inOrder(notificationManager).let {
+            it.verify(notificationManager, times(1)).notify(eq(R.string.stylus_battery_low), any())
+            it.verify(notificationManager, times(1)).cancel(R.string.stylus_battery_low)
+            it.verifyNoMoreInteractions()
+        }
+    }
+
+    @Test
+    fun updateBatteryState_existingNotification_capacityBelowThreshold_updatesNotification() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.15f))
+
+        verify(notificationManager, times(2)).notify(eq(R.string.stylus_battery_low), any())
+        verifyNoMoreInteractions(notificationManager)
+    }
+
+    @Test
+    fun updateBatteryState_capacityAboveThenBelowThreshold_hidesThenShowsNotification() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.5f))
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+
+        inOrder(notificationManager).let {
+            it.verify(notificationManager, times(1)).notify(eq(R.string.stylus_battery_low), any())
+            it.verify(notificationManager, times(1)).cancel(R.string.stylus_battery_low)
+            it.verify(notificationManager, times(1)).notify(eq(R.string.stylus_battery_low), any())
+            it.verifyNoMoreInteractions()
+        }
+    }
+
+    @Test
+    fun updateSuppression_noExistingNotification_cancelsNotification() {
+        stylusUsiPowerUi.updateSuppression(true)
+
+        verify(notificationManager, times(1)).cancel(R.string.stylus_battery_low)
+        verifyNoMoreInteractions(notificationManager)
+    }
+
+    @Test
+    fun updateSuppression_existingNotification_cancelsNotification() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+
+        stylusUsiPowerUi.updateSuppression(true)
+
+        inOrder(notificationManager).let {
+            it.verify(notificationManager, times(1)).notify(eq(R.string.stylus_battery_low), any())
+            it.verify(notificationManager, times(1)).cancel(R.string.stylus_battery_low)
+            it.verifyNoMoreInteractions()
+        }
+    }
+
+    @Test
+    @Ignore("TODO(b/257936830): get bt address once input api available")
+    fun refresh_hasConnectedBluetoothStylus_doesNotNotify() {
+        whenever(inputManager.inputDeviceIds).thenReturn(intArrayOf(0))
+
+        stylusUsiPowerUi.refresh()
+
+        verifyNoMoreInteractions(notificationManager)
+    }
+
+    @Test
+    @Ignore("TODO(b/257936830): get bt address once input api available")
+    fun refresh_hasConnectedBluetoothStylus_existingNotification_cancelsNotification() {
+        stylusUsiPowerUi.updateBatteryState(FixedCapacityBatteryState(0.1f))
+        whenever(inputManager.inputDeviceIds).thenReturn(intArrayOf(0))
+
+        stylusUsiPowerUi.refresh()
+
+        verify(notificationManager).cancel(R.string.stylus_battery_low)
+    }
+
+    class FixedCapacityBatteryState(private val capacity: Float) : BatteryState() {
+        override fun getCapacity() = capacity
+        override fun getStatus() = 0
+        override fun isPresent() = true
+    }
+}