[Status bar refactor] Fix cache removal exception

The old method of removing unused subscriptions from the repo cache was
causing a CME

Test: DemoMobileConnectionsRepositoryTest
Test: MobileConnectionsRepositoryTest
Fixes: 261706421
Change-Id: I4977aa47591a7d4888c176e420cc1a0c7f783591

packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt

@@ -61,7 +61,7 @@ constructor(
 
     private var demoCommandJob: Job? = null
 
-    private val connectionRepoCache = mutableMapOf<Int, DemoMobileConnectionRepository>()
+    private var connectionRepoCache = mutableMapOf<Int, DemoMobileConnectionRepository>()
     private val subscriptionInfoCache = mutableMapOf<Int, SubscriptionModel>()
     val demoModeFinishedEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
 
@@ -76,11 +76,10 @@ constructor(
         // will get garbage collected once their subscribers go away
         val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
 
-        connectionRepoCache.keys.forEach {
-            if (!currentValidSubscriptionIds.contains(it)) {
-                connectionRepoCache.remove(it)
-            }
-        }
+        connectionRepoCache =
+            connectionRepoCache
+                .filter { currentValidSubscriptionIds.contains(it.key) }
+                .toMutableMap()
     }
 
     private fun maybeCreateSubscription(subId: Int) {

packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt

@@ -87,7 +87,7 @@ constructor(
     @Application private val scope: CoroutineScope,
     private val mobileConnectionRepositoryFactory: MobileConnectionRepositoryImpl.Factory
 ) : MobileConnectionsRepository {
-    private val subIdRepositoryCache: MutableMap<Int, MobileConnectionRepository> = mutableMapOf()
+    private var subIdRepositoryCache: MutableMap<Int, MobileConnectionRepository> = mutableMapOf()
 
     /**
      * State flow that emits the set of mobile data subscriptions, each represented by its own
@@ -264,11 +264,10 @@ constructor(
         // will get garbage collected once their subscribers go away
         val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
 
-        subIdRepositoryCache.keys.forEach {
-            if (!currentValidSubscriptionIds.contains(it)) {
-                subIdRepositoryCache.remove(it)
-            }
-        }
+        subIdRepositoryCache =
+            subIdRepositoryCache
+                .filter { currentValidSubscriptionIds.contains(it.key) }
+                .toMutableMap()
     }
 
     private suspend fun fetchSubscriptionsList(): List<SubscriptionInfo> =

packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt

@@ -184,6 +184,25 @@ class DemoMobileConnectionsRepositoryTest : SysuiTestCase() {
             job.cancel()
         }
 
+    /** Regression test for b/261706421 */
+    @Test
+    fun `multiple connections - remove all - does not throw`() =
+        testScope.runTest {
+            var latest: List<SubscriptionModel>? = null
+            val job = underTest.subscriptions.onEach { latest = it }.launchIn(this)
+
+            // Two subscriptions are added
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2, level = 1)
+
+            // Then both are removed by turning off demo mode
+            underTest.stopProcessingCommands()
+
+            assertThat(latest).isEmpty()
+
+            job.cancel()
+        }
+
     @Test
     fun `demo connection - single subscription`() =
         testScope.runTest {

packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt

@@ -52,6 +52,7 @@ import org.junit.After
 import org.junit.Assert.assertThrows
 import org.junit.Before
 import org.junit.Test
+import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
@@ -76,6 +77,24 @@ class MobileConnectionsRepositoryTest : SysuiTestCase() {
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
+        // Set up so the individual connection repositories
+        whenever(telephonyManager.createForSubscriptionId(anyInt())).thenAnswer { invocation ->
+            telephonyManager.also {
+                whenever(telephonyManager.subscriptionId).thenReturn(invocation.getArgument(0))
+            }
+        }
+
+        val connectionFactory: MobileConnectionRepositoryImpl.Factory =
+            MobileConnectionRepositoryImpl.Factory(
+                context = context,
+                telephonyManager = telephonyManager,
+                bgDispatcher = IMMEDIATE,
+                globalSettings = globalSettings,
+                logger = logger,
+                mobileMappingsProxy = mobileMappings,
+                scope = scope,
+            )
+
         underTest =
             MobileConnectionsRepositoryImpl(
                 connectivityManager,
@@ -88,7 +107,7 @@ class MobileConnectionsRepositoryTest : SysuiTestCase() {
                 context,
                 IMMEDIATE,
                 scope,
-                mock(),
+                connectionFactory,
             )
     }
 
@@ -207,6 +226,32 @@ class MobileConnectionsRepositoryTest : SysuiTestCase() {
             job.cancel()
         }
 
+    /** Regression test for b/261706421 */
+    @Test
+    fun testConnectionsCache_clearMultipleSubscriptionsAtOnce_doesNotThrow() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.subscriptions.launchIn(this)
+
+            whenever(subscriptionManager.completeActiveSubscriptionInfoList)
+                .thenReturn(listOf(SUB_1, SUB_2))
+            getSubscriptionCallback().onSubscriptionsChanged()
+
+            // Get repos to trigger caching
+            val repo1 = underTest.getRepoForSubId(SUB_1_ID)
+            val repo2 = underTest.getRepoForSubId(SUB_2_ID)
+
+            assertThat(underTest.getSubIdRepoCache())
+                .containsExactly(SUB_1_ID, repo1, SUB_2_ID, repo2)
+
+            // All subscriptions disappear
+            whenever(subscriptionManager.completeActiveSubscriptionInfoList).thenReturn(listOf())
+            getSubscriptionCallback().onSubscriptionsChanged()
+
+            assertThat(underTest.getSubIdRepoCache()).isEmpty()
+
+            job.cancel()
+        }
+
     @Test
     fun testConnectionRepository_invalidSubId_throws() =
         runBlocking(IMMEDIATE) {