UIEvent for tracking USI battery presence.

Bug: 267816496
Test: StylusManagerTest
Change-Id: I69edb7fc04f4af5ca1768aaf0001f62885bc71d3
diff --git a/packages/SystemUI/src/com/android/systemui/stylus/StylusManager.kt b/packages/SystemUI/src/com/android/systemui/stylus/StylusManager.kt
index 3f54aebf..dee4a6f 100644
--- a/packages/SystemUI/src/com/android/systemui/stylus/StylusManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/stylus/StylusManager.kt
@@ -30,6 +30,9 @@
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.shared.hardware.hasInputDevice
+import com.android.systemui.shared.hardware.isInternalStylusSource
+import com.android.systemui.statusbar.notification.collection.listbuilder.DEBUG
 import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.Executor
 import javax.inject.Inject
@@ -59,8 +62,10 @@
         CopyOnWriteArrayList()
     // This map should only be accessed on the handler
     private val inputDeviceAddressMap: MutableMap<Int, String?> = ArrayMap()
-    // This variable should only be accessed on the handler
+
+    // These variables should only be accessed on the handler
     private var hasStarted: Boolean = false
+    private var isInUsiSession: Boolean = false
 
     /**
      * Starts listening to InputManager InputDevice events. Will also load the InputManager snapshot
@@ -70,6 +75,10 @@
         handler.post {
             if (hasStarted) return@post
             hasStarted = true
+            isInUsiSession =
+                inputManager.hasInputDevice {
+                    it.isInternalStylusSource && isBatteryStateValid(it.batteryState)
+                }
             addExistingStylusToMap()
 
             inputManager.registerInputDeviceListener(this, handler)
@@ -177,7 +186,18 @@
         handler.post {
             if (!hasStarted) return@post
 
-            if (batteryState.isPresent) {
+            if (DEBUG) {
+                Log.d(
+                    TAG,
+                    "onBatteryStateChanged for $deviceId. " +
+                        "batteryState present: ${batteryState.isPresent}, " +
+                        "capacity: ${batteryState.capacity}"
+                )
+            }
+
+            val batteryStateValid = isBatteryStateValid(batteryState)
+            trackAndLogUsiSession(deviceId, batteryStateValid)
+            if (batteryStateValid) {
                 onStylusUsed()
             }
 
@@ -221,6 +241,37 @@
         executeStylusCallbacks { cb -> cb.onStylusFirstUsed() }
     }
 
+    /**
+     * Uses the input device battery state to track whether a current USI session is active. The
+     * InputDevice battery state updates USI battery on USI stylus input, and removes the last-known
+     * USI stylus battery presence after 1 hour of not detecting input. As SysUI and StylusManager
+     * is persistently running, relies on tracking sessions via an in-memory isInUsiSession boolean.
+     */
+    private fun trackAndLogUsiSession(deviceId: Int, batteryStateValid: Boolean) {
+        // TODO(b/268618918) handle cases where an invalid battery callback from a previous stylus
+        //  is sent after the actual valid callback
+        if (batteryStateValid && !isInUsiSession) {
+            if (DEBUG) {
+                Log.d(
+                    TAG,
+                    "USI battery newly present, entering new USI session. Device ID: $deviceId"
+                )
+            }
+            isInUsiSession = true
+            uiEventLogger.log(StylusUiEvent.USI_STYLUS_BATTERY_PRESENCE_FIRST_DETECTED)
+        } else if (!batteryStateValid && isInUsiSession) {
+            if (DEBUG) {
+                Log.d(TAG, "USI battery newly absent, exiting USI session Device ID: $deviceId")
+            }
+            isInUsiSession = false
+            uiEventLogger.log(StylusUiEvent.USI_STYLUS_BATTERY_PRESENCE_REMOVED)
+        }
+    }
+
+    private fun isBatteryStateValid(batteryState: BatteryState): Boolean {
+        return batteryState.isPresent && batteryState.capacity > 0.0f
+    }
+
     private fun executeStylusCallbacks(run: (cb: StylusCallback) -> Unit) {
         stylusCallbacks.forEach(run)
     }
@@ -295,5 +346,6 @@
 
     companion object {
         private val TAG = StylusManager::class.simpleName.orEmpty()
+        private val DEBUG = false
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/stylus/StylusUiEvent.kt b/packages/SystemUI/src/com/android/systemui/stylus/StylusUiEvent.kt
index 99da4ce..e77749b 100644
--- a/packages/SystemUI/src/com/android/systemui/stylus/StylusUiEvent.kt
+++ b/packages/SystemUI/src/com/android/systemui/stylus/StylusUiEvent.kt
@@ -31,7 +31,11 @@
     @UiEvent(doc = "UIEvent for Toast shown when stylus stopped charging")
     STYLUS_STOPPED_CHARGING(1303),
     @UiEvent(doc = "UIEvent for bluetooth stylus connected") BLUETOOTH_STYLUS_CONNECTED(1304),
-    @UiEvent(doc = "UIEvent for bluetooth stylus disconnected") BLUETOOTH_STYLUS_DISCONNECTED(1305);
+    @UiEvent(doc = "UIEvent for bluetooth stylus disconnected") BLUETOOTH_STYLUS_DISCONNECTED(1305),
+    @UiEvent(doc = "UIEvent for start of a USI session via battery presence")
+    USI_STYLUS_BATTERY_PRESENCE_FIRST_DETECTED(1306),
+    @UiEvent(doc = "UIEvent for end of a USI session via battery absence")
+    USI_STYLUS_BATTERY_PRESENCE_REMOVED(1307);
 
     override fun getId() = _id
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt
index 56203d9..6d7941f6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt
@@ -96,6 +96,9 @@
         whenever(stylusDevice.bluetoothAddress).thenReturn(null)
         whenever(btStylusDevice.bluetoothAddress).thenReturn(STYLUS_BT_ADDRESS)
 
+        whenever(stylusDevice.batteryState).thenReturn(batteryState)
+        whenever(batteryState.capacity).thenReturn(0.5f)
+
         whenever(inputManager.getInputDevice(OTHER_DEVICE_ID)).thenReturn(otherDevice)
         whenever(inputManager.getInputDevice(STYLUS_DEVICE_ID)).thenReturn(stylusDevice)
         whenever(inputManager.getInputDevice(BT_STYLUS_DEVICE_ID)).thenReturn(btStylusDevice)
@@ -494,6 +497,47 @@
     }
 
     @Test
+    fun onBatteryStateChanged_batteryPresent_notInUsiSession_logsSessionStart() {
+        whenever(batteryState.isPresent).thenReturn(true)
+
+        stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
+
+        verify(uiEventLogger, times(1))
+            .log(StylusUiEvent.USI_STYLUS_BATTERY_PRESENCE_FIRST_DETECTED)
+    }
+
+    @Test
+    fun onBatteryStateChanged_batteryPresent_inUsiSession_doesNotLogSessionStart() {
+        whenever(batteryState.isPresent).thenReturn(true)
+        stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
+        clearInvocations(uiEventLogger)
+
+        stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
+
+        verify(uiEventLogger, never()).log(any())
+    }
+
+    @Test
+    fun onBatteryStateChanged_batteryAbsent_notInUsiSession_doesNotLogSessionEnd() {
+        whenever(batteryState.isPresent).thenReturn(false)
+
+        stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
+
+        verify(uiEventLogger, never()).log(any())
+    }
+
+    @Test
+    fun onBatteryStateChanged_batteryAbsent_inUsiSession_logSessionEnd() {
+        whenever(batteryState.isPresent).thenReturn(true)
+        stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
+        whenever(batteryState.isPresent).thenReturn(false)
+
+        stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
+
+        verify(uiEventLogger, times(1)).log(StylusUiEvent.USI_STYLUS_BATTERY_PRESENCE_REMOVED)
+    }
+
+    @Test
     fun onBatteryStateChanged_batteryPresent_stylusUsed_doesNotUpdateEverUsedFlag() {
         whenever(inputManager.isStylusEverUsed(mContext)).thenReturn(true)
         whenever(batteryState.isPresent).thenReturn(true)