Merge "Fix potential NPE & clean up some logic" into tm-dev

libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt

@@ -126,7 +126,11 @@ class BubbleVolatileRepository(private val launcherApps: LauncherApps) {
         @UserIdInt userId: Int,
         @UserIdInt parentUserId: Int
     ): Boolean {
-        return entitiesByUser.get(parentUserId).removeIf { b: BubbleEntity -> b.userId == userId }
+        if (entitiesByUser.get(parentUserId) != null) {
+            return entitiesByUser.get(parentUserId).removeIf {
+                b: BubbleEntity -> b.userId == userId }
+        }
+        return false
     }
 
     /**
@@ -141,8 +145,9 @@ class BubbleVolatileRepository(private val launcherApps: LauncherApps) {
             // First check if the user is a parent / top-level user
             val parentUserId = entitiesByUser.keyAt(i)
             if (!activeUsers.contains(parentUserId)) {
-                return removeBubblesForUser(parentUserId, -1)
+                entitiesByUser.remove(parentUserId)
-            } else {
+                return true
+            } else if (entitiesByUser.get(parentUserId) != null) {
                 // Then check if each of the bubbles in the top-level user, still has a valid user
                 // as it could belong to a profile and have a different id from the parent.
                 return entitiesByUser.get(parentUserId).removeIf { b: BubbleEntity ->

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt

@@ -269,6 +269,12 @@ class BubbleVolatileRepositoryTest : ShellTestCase() {
         assertThat(repository.getEntities(user11.identifier).toList())
                 .isEqualTo(listOf(bubble11, bubble12))
     }
+
+    @Test
+    fun testRemoveBubbleForUser_invalidInputDoesntCrash() {
+        repository.removeBubblesForUser(-1, 0)
+        repository.removeBubblesForUser(-1, -1)
+    }
 }
 
 private const val PKG_MESSENGER = "com.example.messenger"