Always update the InputDevice ID in the USI low battery notification

Bug: 290026270
Test: StylusUsiPowerUiTest
Change-Id: I4fae69a083c183235a76da6fe1b0b01a8a71c2e6

packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt

@@ -64,9 +64,9 @@ constructor(
     // These values must only be accessed on the handler.
     private var batteryCapacity = 1.0f
     private var suppressed = false
-    private var inputDeviceId: Int? = null
     private var instanceId: InstanceId? = null
-
+    @VisibleForTesting var inputDeviceId: Int? = null
+      private set
     @VisibleForTesting var instanceIdSequence = InstanceIdSequence(1 shl 13)
 
     fun init() {
@@ -110,10 +110,10 @@ constructor(
 
     fun updateBatteryState(deviceId: Int, batteryState: BatteryState) {
         handler.post updateBattery@{
+            inputDeviceId = deviceId
             if (batteryState.capacity == batteryCapacity || batteryState.capacity <= 0f)
                 return@updateBattery
 
-            inputDeviceId = deviceId
             batteryCapacity = batteryState.capacity
             debugLog {
                 "Updating notification battery state to $batteryCapacity " +

packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt

@@ -125,6 +125,16 @@ class StylusUsiPowerUiTest : SysuiTestCase() {
         verifyNoMoreInteractions(notificationManager)
     }
 
+    @Test
+    fun updateBatteryState_capacitySame_inputDeviceChanges_updatesInputDeviceId() {
+        stylusUsiPowerUi.updateBatteryState(0, FixedCapacityBatteryState(0.1f))
+        stylusUsiPowerUi.updateBatteryState(1, FixedCapacityBatteryState(0.1f))
+
+        assertThat(stylusUsiPowerUi.inputDeviceId).isEqualTo(1)
+        verify(notificationManager, times(1))
+            .notify(eq(R.string.stylus_battery_low_percentage), any())
+    }
+
     @Test
     fun updateBatteryState_existingNotification_capacityAboveThreshold_cancelsNotification() {
         stylusUsiPowerUi.updateBatteryState(0, FixedCapacityBatteryState(0.1f))