diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt index 0d3ba9a78e353..1eee0291cb261 100644 --- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepository.kt +++ b/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) - } else { + entitiesByUser.remove(parentUserId) + 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 -> diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt index 77c70557fb707..9f0d89bc31288 100644 --- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/storage/BubbleVolatileRepositoryTest.kt +++ b/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"