New BroadcastReceiverAsUserFlow

Use callbackFlow to register / unregister receiver to prevent leak.
Before this change, if DisposableBroadcastReceiverAsUser() is removed
from compose before onStop(), leak may happens.

Bug: 308681068
Test: manual - on Settings All Apps, verify app change event collected
Test: unit tests
Change-Id: I73db37d01a0001a2abd1480b44302d61ddbfc695

5 files changed, 224 insertions, 113 deletions

diff --git a/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/common/BroadcastReceiverAsUserFlow.kt b/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/common/BroadcastReceiverAsUserFlow.kt
new file mode 100644
index 000000000000..2c60db4e76c7
--- /dev/null
+++ b/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/common/BroadcastReceiverAsUserFlow.kt
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.spaprivileged.framework.common
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.UserHandle
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.conflate
+import kotlinx.coroutines.flow.flowOn
+
+/**
+ * A [BroadcastReceiver] flow for the given [intentFilter].
+ */
+fun Context.broadcastReceiverAsUserFlow(
+    intentFilter: IntentFilter,
+    userHandle: UserHandle,
+): Flow<Intent> = callbackFlow {
+    val broadcastReceiver = object : BroadcastReceiver() {
+        override fun onReceive(context: Context, intent: Intent) {
+            trySend(intent)
+        }
+    }
+    registerReceiverAsUser(
+        broadcastReceiver,
+        userHandle,
+        intentFilter,
+        null,
+        null,
+        Context.RECEIVER_NOT_EXPORTED,
+    )
+
+    awaitClose { unregisterReceiver(broadcastReceiver) }
+}.conflate().flowOn(Dispatchers.Default)
diff --git a/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUser.kt b/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUser.kt
index ad907cfeb5bf..7d6ee1928111 100644
--- a/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUser.kt
+++ b/packages/SettingsLib/SpaPrivileged/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUser.kt
@@ -17,14 +17,14 @@
 package com.android.settingslib.spaprivileged.framework.compose
 
 import android.content.BroadcastReceiver
-import android.content.Context
 import android.content.Intent
 import android.content.IntentFilter
 import android.os.UserHandle
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.remember
 import androidx.compose.ui.platform.LocalContext
-import com.android.settingslib.spa.framework.compose.LifecycleEffect
+import androidx.compose.ui.platform.LocalLifecycleOwner
+import com.android.settingslib.spa.framework.util.collectLatestWithLifecycle
+import com.android.settingslib.spaprivileged.framework.common.broadcastReceiverAsUserFlow
 
 /**
  * A [BroadcastReceiver] which registered when on start and unregistered when on stop.
@@ -35,27 +35,6 @@ fun DisposableBroadcastReceiverAsUser(
     userHandle: UserHandle,
     onReceive: (Intent) -> Unit,
 ) {
-    val context = LocalContext.current
-    val broadcastReceiver = remember {
-        object : BroadcastReceiver() {
-            override fun onReceive(context: Context, intent: Intent) {
-                onReceive(intent)
-            }
-        }
-    }
-    LifecycleEffect(
-        onStart = {
-            context.registerReceiverAsUser(
-                broadcastReceiver,
-                userHandle,
-                intentFilter,
-                null,
-                null,
-                Context.RECEIVER_NOT_EXPORTED,
-            )
-        },
-        onStop = {
-            context.unregisterReceiver(broadcastReceiver)
-        },
-    )
+    LocalContext.current.broadcastReceiverAsUserFlow(intentFilter, userHandle)
+        .collectLatestWithLifecycle(LocalLifecycleOwner.current, action = onReceive)
 }
diff --git a/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/common/BroadcastReceiverAsUserFlowTest.kt b/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/common/BroadcastReceiverAsUserFlowTest.kt
new file mode 100644
index 000000000000..dfb8e22c7b52
--- /dev/null
+++ b/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/common/BroadcastReceiverAsUserFlowTest.kt
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.spaprivileged.framework.common
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.settingslib.spa.testutils.firstWithTimeoutOrNull
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.isNull
+import org.mockito.kotlin.mock
+
+@RunWith(AndroidJUnit4::class)
+class BroadcastReceiverAsUserFlowTest {
+
+    private var registeredBroadcastReceiver: BroadcastReceiver? = null
+
+    private val context = mock<Context> {
+        on {
+            registerReceiverAsUser(
+                any(),
+                eq(USER_HANDLE),
+                eq(INTENT_FILTER),
+                isNull(),
+                isNull(),
+                eq(Context.RECEIVER_NOT_EXPORTED),
+            )
+        } doAnswer {
+            registeredBroadcastReceiver = it.arguments[0] as BroadcastReceiver
+            null
+        }
+    }
+
+    @Test
+    fun broadcastReceiverAsUserFlow_registered() = runBlocking {
+        val flow = context.broadcastReceiverAsUserFlow(INTENT_FILTER, USER_HANDLE)
+
+        flow.firstWithTimeoutOrNull()
+
+        assertThat(registeredBroadcastReceiver).isNotNull()
+    }
+
+    @Test
+    fun broadcastReceiverAsUserFlow_isCalledOnReceive() = runBlocking {
+        var onReceiveIsCalled = false
+        launch {
+            context.broadcastReceiverAsUserFlow(INTENT_FILTER, USER_HANDLE).first {
+                onReceiveIsCalled = true
+                true
+            }
+        }
+
+        delay(100)
+        registeredBroadcastReceiver!!.onReceive(context, Intent())
+        delay(100)
+
+        assertThat(onReceiveIsCalled).isTrue()
+    }
+
+    private companion object {
+        val USER_HANDLE: UserHandle = UserHandle.of(0)
+
+        val INTENT_FILTER = IntentFilter()
+    }
+}
diff --git a/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUserTest.kt b/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUserTest.kt
index 2c8fb66fab0a..f812f959db32 100644
--- a/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUserTest.kt
+++ b/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/framework/compose/DisposableBroadcastReceiverAsUserTest.kt
@@ -23,38 +23,32 @@ import android.content.IntentFilter
 import android.os.UserHandle
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.platform.LocalLifecycleOwner
 import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.lifecycle.testing.TestLifecycleOwner
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.google.common.truth.Truth.assertThat
-import org.junit.Before
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.runBlocking
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
 import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.isNull
-import org.mockito.kotlin.whenever
+import org.mockito.kotlin.mock
 
 @RunWith(AndroidJUnit4::class)
 class DisposableBroadcastReceiverAsUserTest {
     @get:Rule
     val composeTestRule = createComposeRule()
 
-    @get:Rule
-    val mockito: MockitoRule = MockitoJUnit.rule()
-
-    @Mock
-    private lateinit var context: Context
-
     private var registeredBroadcastReceiver: BroadcastReceiver? = null
 
-    @Before
-    fun setUp() {
-        whenever(
-            context.registerReceiverAsUser(
+    private val context = mock<Context> {
+        on {
+            registerReceiverAsUser(
                 any(),
                 eq(USER_HANDLE),
                 eq(INTENT_FILTER),
@@ -62,7 +56,7 @@ class DisposableBroadcastReceiverAsUserTest {
                 isNull(),
                 eq(Context.RECEIVER_NOT_EXPORTED),
             )
-        ).then {
+        } doAnswer {
             registeredBroadcastReceiver = it.arguments[0] as BroadcastReceiver
             null
         }
@@ -71,7 +65,10 @@ class DisposableBroadcastReceiverAsUserTest {
     @Test
     fun broadcastReceiver_registered() {
         composeTestRule.setContent {
-            CompositionLocalProvider(LocalContext provides context) {
+            CompositionLocalProvider(
+                LocalContext provides context,
+                LocalLifecycleOwner provides TestLifecycleOwner(),
+            ) {
                 DisposableBroadcastReceiverAsUser(INTENT_FILTER, USER_HANDLE) {}
             }
         }
@@ -80,10 +77,13 @@ class DisposableBroadcastReceiverAsUserTest {
     }
 
     @Test
-    fun broadcastReceiver_isCalledOnReceive() {
+    fun broadcastReceiver_isCalledOnReceive() = runBlocking {
         var onReceiveIsCalled = false
         composeTestRule.setContent {
-            CompositionLocalProvider(LocalContext provides context) {
+            CompositionLocalProvider(
+                LocalContext provides context,
+                LocalLifecycleOwner provides TestLifecycleOwner(),
+            ) {
                 DisposableBroadcastReceiverAsUser(INTENT_FILTER, USER_HANDLE) {
                     onReceiveIsCalled = true
                 }
@@ -91,6 +91,7 @@ class DisposableBroadcastReceiverAsUserTest {
         }
 
         registeredBroadcastReceiver!!.onReceive(context, Intent())
+        delay(100)
 
         assertThat(onReceiveIsCalled).isTrue()
     }
diff --git a/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/model/app/AppListRepositoryTest.kt b/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/model/app/AppListRepositoryTest.kt
index 840bca8b14ec..44973a743c76 100644
--- a/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/model/app/AppListRepositoryTest.kt
+++ b/packages/SettingsLib/SpaPrivileged/tests/src/com/android/settingslib/spaprivileged/model/app/AppListRepositoryTest.kt
@@ -19,6 +19,8 @@ package com.android.settingslib.spaprivileged.model.app
 import android.content.Context
 import android.content.pm.ActivityInfo
 import android.content.pm.ApplicationInfo
+import android.content.pm.FakeFeatureFlagsImpl
+import android.content.pm.Flags
 import android.content.pm.PackageManager
 import android.content.pm.PackageManager.ApplicationInfoFlags
 import android.content.pm.PackageManager.ResolveInfoFlags
@@ -29,76 +31,62 @@ import android.os.UserManager
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.internal.R
-import com.android.settingslib.spaprivileged.framework.common.userManager
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.runTest
-import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Spy
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
-import android.content.pm.FakeFeatureFlagsImpl
-import android.content.pm.Flags
 
 @RunWith(AndroidJUnit4::class)
 class AppListRepositoryTest {
-    @get:Rule
-    val mockito: MockitoRule = MockitoJUnit.rule()
-
-    @Spy
-    private val context: Context = ApplicationProvider.getApplicationContext()
-
-    @Mock
-    private lateinit var resources: Resources
-
-    @Mock
-    private lateinit var packageManager: PackageManager
-
-    @Mock
-    private lateinit var userManager: UserManager
-
-    private lateinit var repository: AppListRepository
+    private val resources = mock<Resources> {
+        on { getStringArray(R.array.config_hideWhenDisabled_packageNames) } doReturn emptyArray()
+    }
 
-    @Before
-    fun setUp() {
-        whenever(context.resources).thenReturn(resources)
-        whenever(resources.getStringArray(R.array.config_hideWhenDisabled_packageNames))
-            .thenReturn(emptyArray())
-        whenever(context.packageManager).thenReturn(packageManager)
-        whenever(context.userManager).thenReturn(userManager)
-        whenever(packageManager.getInstalledModules(any())).thenReturn(emptyList())
-        whenever(packageManager.getHomeActivities(any())).thenAnswer {
+    private val packageManager = mock<PackageManager> {
+        on { getInstalledModules(any()) } doReturn emptyList()
+        on { getHomeActivities(any()) } doAnswer {
             @Suppress("UNCHECKED_CAST")
             val resolveInfos = it.arguments[0] as MutableList<ResolveInfo>
             resolveInfos += resolveInfoOf(packageName = HOME_APP.packageName)
             null
         }
-        whenever(
-            packageManager.queryIntentActivitiesAsUser(any(), any<ResolveInfoFlags>(), any<Int>())
-        ).thenReturn(listOf(resolveInfoOf(packageName = IN_LAUNCHER_APP.packageName)))
-        whenever(userManager.getUserInfo(ADMIN_USER_ID)).thenReturn(UserInfo().apply {
+        on { queryIntentActivitiesAsUser(any(), any<ResolveInfoFlags>(), any<Int>()) } doReturn
+            listOf(resolveInfoOf(packageName = IN_LAUNCHER_APP.packageName))
+    }
+
+    private val mockUserManager = mock<UserManager> {
+        on { getUserInfo(ADMIN_USER_ID) } doReturn UserInfo().apply {
             flags = UserInfo.FLAG_ADMIN
-        })
-        whenever(userManager.getProfileIdsWithDisabled(ADMIN_USER_ID))
-            .thenReturn(intArrayOf(ADMIN_USER_ID, MANAGED_PROFILE_USER_ID))
+        }
+        on { getProfileIdsWithDisabled(ADMIN_USER_ID) } doReturn
+            intArrayOf(ADMIN_USER_ID, MANAGED_PROFILE_USER_ID)
+    }
 
-        repository = AppListRepositoryImpl(context)
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+        on { resources } doReturn resources
+        on { packageManager } doReturn packageManager
+        on { getSystemService(UserManager::class.java) } doReturn mockUserManager
     }
 
+    private val repository = AppListRepositoryImpl(context)
+
     private fun mockInstalledApplications(apps: List<ApplicationInfo>, userId: Int) {
-        whenever(
-            packageManager.getInstalledApplicationsAsUser(any<ApplicationInfoFlags>(), eq(userId))
-        ).thenReturn(apps)
+        packageManager.stub {
+            on { getInstalledApplicationsAsUser(any<ApplicationInfoFlags>(), eq(userId)) } doReturn
+                apps
+        }
     }
 
     @Test
@@ -135,13 +123,13 @@ class AppListRepositoryTest {
         )
 
         assertThat(appList).containsExactly(NORMAL_APP)
-        argumentCaptor<ApplicationInfoFlags> {
+        val flags = argumentCaptor<ApplicationInfoFlags> {
             verify(packageManager).getInstalledApplicationsAsUser(capture(), eq(ADMIN_USER_ID))
-            assertThat(firstValue.value).isEqualTo(
-                PackageManager.MATCH_DISABLED_COMPONENTS or
-                    PackageManager.MATCH_DISABLED_UNTIL_USED_COMPONENTS
-            )
-        }
+        }.firstValue
+        assertThat(flags.value).isEqualTo(
+            PackageManager.MATCH_DISABLED_COMPONENTS or
+                PackageManager.MATCH_DISABLED_UNTIL_USED_COMPONENTS
+        )
     }
 
     @Test
@@ -154,11 +142,10 @@ class AppListRepositoryTest {
         )
 
         assertThat(appList).containsExactly(NORMAL_APP)
-        argumentCaptor<ApplicationInfoFlags> {
+        val flags = argumentCaptor<ApplicationInfoFlags> {
             verify(packageManager).getInstalledApplicationsAsUser(capture(), eq(ADMIN_USER_ID))
-            assertThat(firstValue.value and PackageManager.MATCH_ANY_USER.toLong())
-                .isGreaterThan(0L)
-        }
+        }.firstValue
+        assertThat(flags.value and PackageManager.MATCH_ANY_USER.toLong()).isGreaterThan(0L)
     }
 
     @Test
@@ -278,14 +265,14 @@ class AppListRepositoryTest {
         val appList = repository.loadApps(userId = ADMIN_USER_ID)
 
         assertThat(appList).containsExactly(NORMAL_APP, ARCHIVED_APP)
-        argumentCaptor<ApplicationInfoFlags> {
+        val flags = argumentCaptor<ApplicationInfoFlags> {
             verify(packageManager).getInstalledApplicationsAsUser(capture(), eq(ADMIN_USER_ID))
-            assertThat(firstValue.value).isEqualTo(
-                (PackageManager.MATCH_DISABLED_COMPONENTS or
-                    PackageManager.MATCH_DISABLED_UNTIL_USED_COMPONENTS).toLong() or
-                    PackageManager.MATCH_ARCHIVED_PACKAGES
-            )
-        }
+        }.firstValue
+        assertThat(flags.value).isEqualTo(
+            (PackageManager.MATCH_DISABLED_COMPONENTS or
+                PackageManager.MATCH_DISABLED_UNTIL_USED_COMPONENTS).toLong() or
+                PackageManager.MATCH_ARCHIVED_PACKAGES
+        )
     }
 
     @Test
@@ -294,13 +281,13 @@ class AppListRepositoryTest {
         val appList = repository.loadApps(userId = ADMIN_USER_ID)
 
         assertThat(appList).containsExactly(NORMAL_APP)
-        argumentCaptor<ApplicationInfoFlags> {
+        val flags = argumentCaptor<ApplicationInfoFlags> {
             verify(packageManager).getInstalledApplicationsAsUser(capture(), eq(ADMIN_USER_ID))
-            assertThat(firstValue.value).isEqualTo(
-                PackageManager.MATCH_DISABLED_COMPONENTS or
-                    PackageManager.MATCH_DISABLED_UNTIL_USED_COMPONENTS
-            )
-        }
+        }.firstValue
+        assertThat(flags.value).isEqualTo(
+            PackageManager.MATCH_DISABLED_COMPONENTS or
+                PackageManager.MATCH_DISABLED_UNTIL_USED_COMPONENTS
+        )
     }
 
     @Test