Merge "Partial Screen Sharing: use correct icon badges in recent apps selector" into main

10 files changed, 165 insertions, 54 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
index d2471225a093..f08bc17c4f23 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
@@ -25,8 +25,8 @@ import com.android.launcher3.icons.IconFactory
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.mediaprojection.appselector.data.ActivityTaskManagerLabelLoader
 import com.android.systemui.mediaprojection.appselector.data.ActivityTaskManagerThumbnailLoader
-import com.android.systemui.mediaprojection.appselector.data.AppIconLoader
-import com.android.systemui.mediaprojection.appselector.data.IconLoaderLibAppIconLoader
+import com.android.systemui.mediaprojection.appselector.data.BasicAppIconLoader
+import com.android.systemui.mediaprojection.appselector.data.BasicPackageManagerAppIconLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskLabelLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskListProvider
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskThumbnailLoader
@@ -102,7 +102,7 @@ interface MediaProjectionAppSelectorModule {
 
     @Binds
     @MediaProjectionAppSelectorScope
-    fun bindAppIconLoader(impl: IconLoaderLibAppIconLoader): AppIconLoader
+    fun bindAppIconLoader(impl: BasicPackageManagerAppIconLoader): BasicAppIconLoader
 
     @Binds
     @IntoSet
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BadgedAppIconLoader.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BadgedAppIconLoader.kt
new file mode 100644
index 000000000000..ca5b5f842b25
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BadgedAppIconLoader.kt
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 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.systemui.mediaprojection.appselector.data
+
+import android.content.ComponentName
+import android.content.Context
+import android.graphics.drawable.Drawable
+import android.os.UserHandle
+import com.android.launcher3.icons.BaseIconFactory
+import com.android.launcher3.icons.IconFactory
+import com.android.launcher3.util.UserIconInfo
+import com.android.systemui.dagger.qualifiers.Background
+import javax.inject.Inject
+import javax.inject.Provider
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+class BadgedAppIconLoader
+@Inject
+constructor(
+    private val basicAppIconLoader: BasicAppIconLoader,
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
+    private val context: Context,
+    private val iconFactoryProvider: Provider<IconFactory>,
+) {
+
+    suspend fun loadIcon(
+        userId: Int,
+        userType: RecentTask.UserType,
+        componentName: ComponentName
+    ): Drawable? =
+        withContext(backgroundDispatcher) {
+            iconFactoryProvider.get().use<IconFactory, Drawable?> { iconFactory ->
+                val icon =
+                    basicAppIconLoader.loadIcon(userId, componentName) ?: return@withContext null
+                val userHandler = UserHandle.of(userId)
+                val iconType = getIconType(userType)
+                val options =
+                    BaseIconFactory.IconOptions().apply {
+                        setUser(UserIconInfo(userHandler, iconType))
+                    }
+                val badgedIcon = iconFactory.createBadgedIconBitmap(icon, options)
+                badgedIcon.newIcon(context)
+            }
+        }
+
+    private fun getIconType(userType: RecentTask.UserType): Int =
+        when (userType) {
+            RecentTask.UserType.CLONED -> UserIconInfo.TYPE_CLONED
+            RecentTask.UserType.WORK -> UserIconInfo.TYPE_WORK
+            RecentTask.UserType.PRIVATE -> UserIconInfo.TYPE_PRIVATE
+            RecentTask.UserType.STANDARD -> UserIconInfo.TYPE_MAIN
+        }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/AppIconLoader.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BasicAppIconLoader.kt
index b85d6285c35b..03f6f015300f 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/AppIconLoader.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BasicAppIconLoader.kt
@@ -17,45 +17,29 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.content.ComponentName
-import android.content.Context
 import android.content.pm.PackageManager
 import android.graphics.drawable.Drawable
-import android.os.UserHandle
-import com.android.launcher3.icons.BaseIconFactory.IconOptions
-import com.android.launcher3.icons.IconFactory
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.shared.system.PackageManagerWrapper
 import javax.inject.Inject
-import javax.inject.Provider
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.withContext
 
-interface AppIconLoader {
+interface BasicAppIconLoader {
     suspend fun loadIcon(userId: Int, component: ComponentName): Drawable?
 }
 
-class IconLoaderLibAppIconLoader
+class BasicPackageManagerAppIconLoader
 @Inject
 constructor(
     @Background private val backgroundDispatcher: CoroutineDispatcher,
-    private val context: Context,
     // Use wrapper to access hidden API that allows to get ActivityInfo for any user id
     private val packageManagerWrapper: PackageManagerWrapper,
     private val packageManager: PackageManager,
-    private val iconFactoryProvider: Provider<IconFactory>
-) : AppIconLoader {
+) : BasicAppIconLoader {
 
     override suspend fun loadIcon(userId: Int, component: ComponentName): Drawable? =
         withContext(backgroundDispatcher) {
-            iconFactoryProvider.get().use<IconFactory, Drawable?> { iconFactory ->
-                val activityInfo =
-                    packageManagerWrapper.getActivityInfo(component, userId)
-                        ?: return@withContext null
-                val icon = activityInfo.loadIcon(packageManager) ?: return@withContext null
-                val userHandler = UserHandle.of(userId)
-                val options = IconOptions().apply { setUser(userHandler) }
-                val badgedIcon = iconFactory.createBadgedIconBitmap(icon, options)
-                badgedIcon.newIcon(context)
-            }
+            packageManagerWrapper.getActivityInfo(component, userId)?.loadIcon(packageManager)
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
index e9b458271ef7..3e9b546d58c9 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
@@ -28,4 +28,12 @@ data class RecentTask(
     val baseIntentComponent: ComponentName?,
     @ColorInt val colorBackground: Int?,
     val isForegroundTask: Boolean,
-)
+    val userType: UserType,
+) {
+    enum class UserType {
+        STANDARD,
+        WORK,
+        PRIVATE,
+        CLONED
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
index 5dde14bf0867..a6049c8b556d 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
@@ -17,6 +17,8 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.app.ActivityManager.RECENT_IGNORE_UNAVAILABLE
+import android.content.pm.UserInfo
+import android.os.UserManager
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.kotlin.getOrNull
@@ -41,7 +43,8 @@ constructor(
     @Background private val coroutineDispatcher: CoroutineDispatcher,
     @Background private val backgroundExecutor: Executor,
     private val recentTasks: Optional<RecentTasks>,
-    private val userTracker: UserTracker
+    private val userTracker: UserTracker,
+    private val userManager: UserManager,
 ) : RecentTaskListProvider {
 
     private val recents by lazy { recentTasks.getOrNull() }
@@ -65,7 +68,8 @@ constructor(
                         it.topActivity,
                         it.baseIntent?.component,
                         it.taskDescription?.backgroundColor,
-                        isForegroundTask = it.taskId in foregroundTaskIds && it.isVisible
+                        isForegroundTask = it.taskId in foregroundTaskIds && it.isVisible,
+                        userType = userManager.getUserInfo(it.userId).toUserType(),
                     )
                 }
         }
@@ -81,4 +85,15 @@ constructor(
                 continuation.resume(tasks)
             }
         }
+
+    private fun UserInfo.toUserType(): RecentTask.UserType =
+        if (isCloneProfile) {
+            RecentTask.UserType.CLONED
+        } else if (isManagedProfile) {
+            RecentTask.UserType.WORK
+        } else if (isPrivateProfile) {
+            RecentTask.UserType.PRIVATE
+        } else {
+            RecentTask.UserType.STANDARD
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt
index 3fe040a0d715..3b84d2c53a2b 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt
@@ -21,12 +21,12 @@ import android.view.View
 import android.view.ViewGroup
 import android.widget.ImageView
 import androidx.recyclerview.widget.RecyclerView.ViewHolder
-import com.android.systemui.res.R
 import com.android.systemui.mediaprojection.appselector.MediaProjectionAppSelector
-import com.android.systemui.mediaprojection.appselector.data.AppIconLoader
+import com.android.systemui.mediaprojection.appselector.data.BadgedAppIconLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTask
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskLabelLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskThumbnailLoader
+import com.android.systemui.res.R
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
@@ -39,7 +39,7 @@ class RecentTaskViewHolder
 @AssistedInject
 constructor(
     @Assisted private val root: ViewGroup,
-    private val iconLoader: AppIconLoader,
+    private val iconLoader: BadgedAppIconLoader,
     private val thumbnailLoader: RecentTaskThumbnailLoader,
     private val labelLoader: RecentTaskLabelLoader,
     private val taskViewSizeProvider: TaskPreviewSizeProvider,
@@ -63,7 +63,7 @@ constructor(
             scope.launch {
                 task.baseIntentComponent?.let { component ->
                     launch {
-                        val icon = iconLoader.loadIcon(task.userId, component)
+                        val icon = iconLoader.loadIcon(task.userId, task.userType, component)
                         iconView.setImageDrawable(icon)
                     }
                     launch {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
index 44798ea99bee..8b79fa45b8ba 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
@@ -257,6 +257,7 @@ class MediaProjectionAppSelectorControllerTest : SysuiTestCase() {
             userId = userId,
             colorBackground = 0,
             isForegroundTask = isForegroundTask,
+            userType = RecentTask.UserType.STANDARD,
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/IconLoaderLibAppIconLoaderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/BasicPackageManagerAppIconLoaderTest.kt
index 9b346d0120ef..fa1c8f8ea2c7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/IconLoaderLibAppIconLoaderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/BasicPackageManagerAppIconLoaderTest.kt
@@ -20,15 +20,10 @@ import android.content.ComponentName
 import android.content.pm.ActivityInfo
 import android.content.pm.PackageManager
 import android.graphics.Bitmap
-import android.graphics.drawable.Drawable
 import androidx.test.filters.SmallTest
-import com.android.launcher3.icons.BitmapInfo
 import com.android.launcher3.icons.FastBitmapDrawable
-import com.android.launcher3.icons.IconFactory
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.shared.system.PackageManagerWrapper
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
@@ -40,20 +35,17 @@ import org.junit.runners.JUnit4
 
 @SmallTest
 @RunWith(JUnit4::class)
-class IconLoaderLibAppIconLoaderTest : SysuiTestCase() {
+class BasicPackageManagerAppIconLoaderTest : SysuiTestCase() {
 
-    private val iconFactory: IconFactory = mock()
     private val packageManagerWrapper: PackageManagerWrapper = mock()
     private val packageManager: PackageManager = mock()
     private val dispatcher = Dispatchers.Unconfined
 
     private val appIconLoader =
-        IconLoaderLibAppIconLoader(
+        BasicPackageManagerAppIconLoader(
             backgroundDispatcher = dispatcher,
-            context = context,
             packageManagerWrapper = packageManagerWrapper,
             packageManager = packageManager,
-            iconFactoryProvider = { iconFactory }
         )
 
     @Test
@@ -70,12 +62,7 @@ class IconLoaderLibAppIconLoaderTest : SysuiTestCase() {
     private fun givenIcon(component: ComponentName, userId: Int, icon: FastBitmapDrawable) {
         val activityInfo = mock<ActivityInfo>()
         whenever(packageManagerWrapper.getActivityInfo(component, userId)).thenReturn(activityInfo)
-        val rawIcon = mock<Drawable>()
-        whenever(activityInfo.loadIcon(packageManager)).thenReturn(rawIcon)
-
-        val bitmapInfo = mock<BitmapInfo>()
-        whenever(iconFactory.createBadgedIconBitmap(eq(rawIcon), any())).thenReturn(bitmapInfo)
-        whenever(bitmapInfo.newIcon(context)).thenReturn(icon)
+        whenever(activityInfo.loadIcon(packageManager)).thenReturn(icon)
     }
 
     private fun createIcon(): FastBitmapDrawable =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
index b593def283ae..dd621129ad9a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
@@ -1,9 +1,15 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.app.ActivityManager.RecentTaskInfo
+import android.content.pm.UserInfo
+import android.os.UserManager
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.CLONED
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.PRIVATE
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.STANDARD
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.WORK
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
@@ -17,6 +23,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
 
 @RunWith(AndroidTestingRunner::class)
 @SmallTest
@@ -25,12 +32,16 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
     private val dispatcher = Dispatchers.Unconfined
     private val recentTasks: RecentTasks = mock()
     private val userTracker: UserTracker = mock()
+    private val userManager: UserManager = mock {
+        whenever(getUserInfo(anyInt())).thenReturn(mock())
+    }
     private val recentTaskListProvider =
         ShellRecentTaskListProvider(
             dispatcher,
             Runnable::run,
             Optional.of(recentTasks),
-            userTracker
+            userTracker,
+            userManager,
         )
 
     @Test
@@ -147,6 +158,22 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
             .inOrder()
     }
 
+    @Test
+    fun loadRecentTasks_assignsCorrectUserType() {
+        givenRecentTasks(
+            createSingleTask(taskId = 1, userId = 10, userType = STANDARD),
+            createSingleTask(taskId = 2, userId = 20, userType = WORK),
+            createSingleTask(taskId = 3, userId = 30, userType = CLONED),
+            createSingleTask(taskId = 4, userId = 40, userType = PRIVATE),
+        )
+
+        val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
+
+        assertThat(result.map { it.userType })
+            .containsExactly(STANDARD, WORK, CLONED, PRIVATE)
+            .inOrder()
+    }
+
     @Suppress("UNCHECKED_CAST")
     private fun givenRecentTasks(vararg tasks: GroupedRecentTaskInfo) {
         whenever(recentTasks.getRecentTasks(any(), any(), any(), any(), any())).thenAnswer {
@@ -155,7 +182,10 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
         }
     }
 
-    private fun createRecentTask(taskId: Int): RecentTask =
+    private fun createRecentTask(
+        taskId: Int,
+        userType: RecentTask.UserType = STANDARD
+    ): RecentTask =
         RecentTask(
             taskId = taskId,
             displayId = 0,
@@ -164,25 +194,42 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
             baseIntentComponent = null,
             colorBackground = null,
             isForegroundTask = false,
+            userType = userType,
         )
 
-    private fun createSingleTask(taskId: Int, isVisible: Boolean = false): GroupedRecentTaskInfo =
-        GroupedRecentTaskInfo.forSingleTask(createTaskInfo(taskId, isVisible))
+    private fun createSingleTask(
+        taskId: Int,
+        userId: Int = 0,
+        isVisible: Boolean = false,
+        userType: RecentTask.UserType = STANDARD,
+    ): GroupedRecentTaskInfo {
+        val userInfo =
+            mock<UserInfo> {
+                whenever(isCloneProfile).thenReturn(userType == CLONED)
+                whenever(isManagedProfile).thenReturn(userType == WORK)
+                whenever(isPrivateProfile).thenReturn(userType == PRIVATE)
+            }
+        whenever(userManager.getUserInfo(userId)).thenReturn(userInfo)
+        return GroupedRecentTaskInfo.forSingleTask(createTaskInfo(taskId, userId, isVisible))
+    }
 
     private fun createTaskPair(
         taskId1: Int,
+        userId1: Int = 0,
         taskId2: Int,
+        userId2: Int = 0,
         isVisible: Boolean = false
     ): GroupedRecentTaskInfo =
         GroupedRecentTaskInfo.forSplitTasks(
-            createTaskInfo(taskId1, isVisible),
-            createTaskInfo(taskId2, isVisible),
+            createTaskInfo(taskId1, userId1, isVisible),
+            createTaskInfo(taskId2, userId2, isVisible),
             null
         )
 
-    private fun createTaskInfo(taskId: Int, isVisible: Boolean = false) =
+    private fun createTaskInfo(taskId: Int, userId: Int, isVisible: Boolean = false) =
         RecentTaskInfo().apply {
             this.taskId = taskId
             this.isVisible = isVisible
+            this.userId = userId
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt
index ac4107359dbc..a84008b0353c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt
@@ -56,7 +56,8 @@ class MediaProjectionRecentsViewControllerTest : SysuiTestCase() {
             topActivityComponent = null,
             baseIntentComponent = null,
             colorBackground = null,
-            isForegroundTask = false
+            isForegroundTask = false,
+            userType = RecentTask.UserType.STANDARD,
         )
 
     private val taskView =