2 files changed, 22 insertions, 1 deletions
diff --git a/java/src/com/android/intentresolver/v2/data/repository/UserRepository.kt b/java/src/com/android/intentresolver/v2/data/repository/UserRepository.kt
index 91ad6409..b57609e5 100644
--- a/java/src/com/android/intentresolver/v2/data/repository/UserRepository.kt
+++ b/java/src/com/android/intentresolver/v2/data/repository/UserRepository.kt
@@ -158,11 +158,14 @@ constructor(
         }
     }
 
+    private fun List<UserWithState>.update(handle: UserHandle, user: UserWithState) =
+        filter { it.user.id != handle.identifier } + user
+
     private fun handleAvailability(event: UserEvent, current: UserStates): UserStates {
         val userEntry =
             current.firstOrNull { it.user.id == event.user.identifier }
                 ?: throw UserStateException("User was not present in the map", event)
-        return current + userEntry.copy(available = !event.quietMode)
+        return current.update(event.user, userEntry.copy(available = !event.quietMode))
     }
 
     private fun handleProfileRemoved(event: UserEvent, current: UserStates): UserStates {
diff --git a/tests/unit/src/com/android/intentresolver/v2/data/repository/UserRepositoryImplTest.kt b/tests/unit/src/com/android/intentresolver/v2/data/repository/UserRepositoryImplTest.kt
index 6c61dfd6..16e8c9bb 100644
--- a/tests/unit/src/com/android/intentresolver/v2/data/repository/UserRepositoryImplTest.kt
+++ b/tests/unit/src/com/android/intentresolver/v2/data/repository/UserRepositoryImplTest.kt
@@ -79,6 +79,24 @@ internal class UserRepositoryImplTest {
     }
 
     @Test
+    fun onHandleAvailabilityChange_userStateMaintained() = runTest {
+        val repo = createUserRepository(userManager)
+        val private = userState.createProfile(ProfileType.PRIVATE)
+        val privateUser = User(private.identifier, Role.PRIVATE)
+
+        val users by collectLastValue(repo.users)
+
+        repo.requestState(privateUser, false)
+        repo.requestState(privateUser, true)
+
+        assertWithMessage("users.size")
+                .that(users?.size ?: 0).isEqualTo(2) // personal + private
+
+        assertWithMessage("No duplicate IDs")
+                .that(users?.count { it.id == private.identifier }).isEqualTo(1)
+    }
+
+    @Test
     fun requestState() = runTest {
         val repo = createUserRepository(userManager)
         val work = userState.createProfile(ProfileType.WORK)