diff --git a/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt b/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
index 3e1c13c1cba82..c1ac800b8159c 100644
--- a/packages/SystemUI/src/com/android/systemui/stylus/StylusUsiPowerUI.kt
+++ b/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 " +
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 90821bdef0bef..d212c026d66eb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusUsiPowerUiTest.kt
+++ b/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))