Filter out work profile widgets when not allowed by device policy

Existing work profile widgets should be filtered out if keyguard widgets
are disabled on the work profile but not on the main user by device policy manager.

Bug: b/323196422
Test: atest CommunalInteractorTest
Test: atest CommunalSettingsRepositoryImplTest
Test: verify by toggling the setting in TestDPC
Flag: ACONFIG com.android.systemui.communal_hub TEAMFOOD
Change-Id: I0ac2d766e4f1942a99f888f1ba8aff43cd700432
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt
index f71121c..ce7b60e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSettingsRepositoryImplTest.kt
@@ -23,6 +23,7 @@
 import android.appwidget.AppWidgetProviderInfo
 import android.content.Intent
 import android.content.pm.UserInfo
+import android.os.UserManager.USER_TYPE_PROFILE_MANAGED
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.provider.Settings
@@ -59,6 +60,7 @@
         kosmos.fakeFeatureFlagsClassic.set(COMMUNAL_SERVICE_ENABLED, true)
         setKeyguardFeaturesDisabled(PRIMARY_USER, KEYGUARD_DISABLE_FEATURES_NONE)
         setKeyguardFeaturesDisabled(SECONDARY_USER, KEYGUARD_DISABLE_FEATURES_NONE)
+        setKeyguardFeaturesDisabled(WORK_PROFILE, KEYGUARD_DISABLE_FEATURES_NONE)
         underTest = kosmos.communalSettingsRepository
     }
 
@@ -133,6 +135,30 @@
 
     @EnableFlags(FLAG_COMMUNAL_HUB)
     @Test
+    fun widgetsAllowedForWorkProfile_isFalse_whenDisallowedByDevicePolicy() =
+        testScope.runTest {
+            val widgetsAllowedForWorkProfile by
+                collectLastValue(underTest.getAllowedByDevicePolicy(WORK_PROFILE))
+            assertThat(widgetsAllowedForWorkProfile).isTrue()
+
+            setKeyguardFeaturesDisabled(WORK_PROFILE, KEYGUARD_DISABLE_WIDGETS_ALL)
+            assertThat(widgetsAllowedForWorkProfile).isFalse()
+        }
+
+    @EnableFlags(FLAG_COMMUNAL_HUB)
+    @Test
+    fun hubIsEnabled_whenDisallowedByDevicePolicyForWorkProfile() =
+        testScope.runTest {
+            val enabledStateForPrimaryUser by
+                collectLastValue(underTest.getEnabledState(PRIMARY_USER))
+            assertThat(enabledStateForPrimaryUser?.enabled).isTrue()
+
+            setKeyguardFeaturesDisabled(WORK_PROFILE, KEYGUARD_DISABLE_WIDGETS_ALL)
+            assertThat(enabledStateForPrimaryUser?.enabled).isTrue()
+        }
+
+    @EnableFlags(FLAG_COMMUNAL_HUB)
+    @Test
     fun hubIsDisabledByUserAndDevicePolicy() =
         testScope.runTest {
             val enabledState by collectLastValue(underTest.getEnabledState(PRIMARY_USER))
@@ -189,5 +215,13 @@
         val PRIMARY_USER =
             UserInfo(/* id= */ 0, /* name= */ "primary user", /* flags= */ UserInfo.FLAG_MAIN)
         val SECONDARY_USER = UserInfo(/* id= */ 1, /* name= */ "secondary user", /* flags= */ 0)
+        val WORK_PROFILE =
+            UserInfo(
+                10,
+                "work",
+                /* iconPath= */ "",
+                /* flags= */ 0,
+                USER_TYPE_PROFILE_MANAGED,
+            )
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
index e7ccde2..f21e969 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
@@ -17,6 +17,8 @@
 
 package com.android.systemui.communal.domain.interactor
 
+import android.app.admin.DevicePolicyManager
+import android.app.admin.devicePolicyManager
 import android.app.smartspace.SmartspaceTarget
 import android.appwidget.AppWidgetProviderInfo
 import android.content.Intent
@@ -32,6 +34,7 @@
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.systemui.Flags.FLAG_COMMUNAL_HUB
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.broadcast.broadcastDispatcher
 import com.android.systemui.communal.data.repository.CommunalSettingsRepositoryImpl
 import com.android.systemui.communal.data.repository.FakeCommunalMediaRepository
 import com.android.systemui.communal.data.repository.FakeCommunalPrefsRepository
@@ -71,6 +74,7 @@
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
@@ -929,7 +933,6 @@
             keyguardRepository.setKeyguardShowing(true)
             keyguardRepository.setKeyguardOccluded(false)
             tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
-            userRepository.setSelectedUserInfo(mainUser)
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
@@ -937,6 +940,7 @@
                 userInfos = userInfos,
                 selectedUserIndex = 0,
             )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
             runCurrent()
 
             // Widgets available.
@@ -955,7 +959,6 @@
                 AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD,
                 mainUser.id
             )
-            runCurrent()
 
             // Only the keyguard widget is enabled.
             assertThat(widgetContent).hasSize(3)
@@ -974,7 +977,6 @@
             keyguardRepository.setKeyguardShowing(true)
             keyguardRepository.setKeyguardOccluded(false)
             tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
-            userRepository.setSelectedUserInfo(mainUser)
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
@@ -982,6 +984,7 @@
                 userInfos = userInfos,
                 selectedUserIndex = 0,
             )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
             runCurrent()
 
             // Widgets available.
@@ -1001,7 +1004,6 @@
                     AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN,
                 mainUser.id
             )
-            runCurrent()
 
             // All widgets are enabled.
             assertThat(widgetContent).hasSize(3)
@@ -1011,6 +1013,79 @@
             }
         }
 
+    @Test
+    fun filterWidgets_whenDisallowedByDevicePolicyForWorkProfile() =
+        testScope.runTest {
+            // Keyguard showing, and tutorial completed.
+            keyguardRepository.setKeyguardShowing(true)
+            keyguardRepository.setKeyguardOccluded(false)
+            tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
+
+            val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
+            userRepository.setUserInfos(userInfos)
+            userTracker.set(
+                userInfos = userInfos,
+                selectedUserIndex = 0,
+            )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
+            runCurrent()
+
+            val widgetContent by collectLastValue(underTest.widgetContent)
+            // Given three widgets, and one of them is associated with work profile.
+            val widget1 = createWidgetForUser(1, USER_INFO_WORK.id)
+            val widget2 = createWidgetForUser(2, MAIN_USER_INFO.id)
+            val widget3 = createWidgetForUser(3, MAIN_USER_INFO.id)
+            val widgets = listOf(widget1, widget2, widget3)
+            widgetRepository.setCommunalWidgets(widgets)
+
+            setKeyguardFeaturesDisabled(
+                USER_INFO_WORK,
+                DevicePolicyManager.KEYGUARD_DISABLE_WIDGETS_ALL
+            )
+
+            // Widget under work profile is filtered out and the remaining two link to main user id.
+            assertThat(widgetContent).hasSize(2)
+            widgetContent!!.forEach { model ->
+                assertThat(model.providerInfo.profile?.identifier).isEqualTo(MAIN_USER_INFO.id)
+            }
+        }
+
+    @Test
+    fun filterWidgets_whenAllowedByDevicePolicyForWorkProfile() =
+        testScope.runTest {
+            // Keyguard showing, and tutorial completed.
+            keyguardRepository.setKeyguardShowing(true)
+            keyguardRepository.setKeyguardOccluded(false)
+            tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
+
+            val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
+            userRepository.setUserInfos(userInfos)
+            userTracker.set(
+                userInfos = userInfos,
+                selectedUserIndex = 0,
+            )
+            userRepository.setSelectedUserInfo(MAIN_USER_INFO)
+            runCurrent()
+
+            val widgetContent by collectLastValue(underTest.widgetContent)
+            // Given three widgets, and one of them is associated with work profile.
+            val widget1 = createWidgetForUser(1, USER_INFO_WORK.id)
+            val widget2 = createWidgetForUser(2, MAIN_USER_INFO.id)
+            val widget3 = createWidgetForUser(3, MAIN_USER_INFO.id)
+            val widgets = listOf(widget1, widget2, widget3)
+            widgetRepository.setCommunalWidgets(widgets)
+
+            setKeyguardFeaturesDisabled(
+                USER_INFO_WORK,
+                DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_NONE
+            )
+
+            // Widget under work profile is available.
+            assertThat(widgetContent).hasSize(3)
+            assertThat(widgetContent!![0].providerInfo.profile?.identifier)
+                .isEqualTo(USER_INFO_WORK.id)
+        }
+
     private fun smartspaceTimer(id: String, timestamp: Long = 0L): SmartspaceTarget {
         val timer = mock(SmartspaceTarget::class.java)
         whenever(timer.smartspaceTargetId).thenReturn(id)
@@ -1020,6 +1095,15 @@
         return timer
     }
 
+    private fun setKeyguardFeaturesDisabled(user: UserInfo, disabledFlags: Int) {
+        whenever(kosmos.devicePolicyManager.getKeyguardDisabledFeatures(nullable(), eq(user.id)))
+            .thenReturn(disabledFlags)
+        kosmos.broadcastDispatcher.sendIntentToMatchingReceiversOnly(
+            context,
+            Intent(DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED),
+        )
+    }
+
     private fun createWidgetForUser(appWidgetId: Int, userId: Int): CommunalWidgetContentModel =
         mock<CommunalWidgetContentModel> {
             whenever(this.appWidgetId).thenReturn(appWidgetId)
@@ -1044,6 +1128,13 @@
 
     private companion object {
         val MAIN_USER_INFO = UserInfo(0, "primary", UserInfo.FLAG_MAIN)
-        val USER_INFO_WORK = UserInfo(10, "work", UserInfo.FLAG_PROFILE)
+        val USER_INFO_WORK =
+            UserInfo(
+                10,
+                "work",
+                /* iconPath= */ "",
+                /* flags= */ 0,
+                UserManager.USER_TYPE_PROFILE_MANAGED,
+            )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt
index c724244..9debe0e 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSettingsRepository.kt
@@ -57,6 +57,9 @@
      * Settings.
      */
     fun getWidgetCategories(user: UserInfo): Flow<CommunalWidgetCategories>
+
+    /** Keyguard widgets enabled state by Device Policy Manager for the specified user. */
+    fun getAllowedByDevicePolicy(user: UserInfo): Flow<Boolean>
 }
 
 @SysUISingleton
@@ -115,6 +118,16 @@
             }
             .flowOn(bgDispatcher)
 
+    override fun getAllowedByDevicePolicy(user: UserInfo): Flow<Boolean> =
+        broadcastDispatcher
+            .broadcastFlow(
+                filter =
+                    IntentFilter(DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED),
+                user = user.userHandle
+            )
+            .emitOnStart()
+            .map { devicePolicyManager.areKeyguardWidgetsAllowed(user.id) }
+
     private fun getEnabledByUser(user: UserInfo): Flow<Boolean> =
         secureSettings
             .observerFlow(userId = user.id, names = arrayOf(Settings.Secure.GLANCEABLE_HUB_ENABLED))
@@ -128,16 +141,6 @@
                 ) == 1
             }
 
-    private fun getAllowedByDevicePolicy(user: UserInfo): Flow<Boolean> =
-        broadcastDispatcher
-            .broadcastFlow(
-                filter =
-                    IntentFilter(DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED),
-                user = user.userHandle
-            )
-            .emitOnStart()
-            .map { devicePolicyManager.areKeyguardWidgetsAllowed(user.id) }
-
     companion object {
         const val GLANCEABLE_HUB_CONTENT_SETTING = "glanceable_hub_content_setting"
         private const val ENABLED_SETTING_DEFAULT = 1
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
index ac8e5e8..373e1c9 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
@@ -99,7 +99,7 @@
     mediaRepository: CommunalMediaRepository,
     smartspaceRepository: SmartspaceRepository,
     keyguardInteractor: KeyguardInteractor,
-    communalSettingsInteractor: CommunalSettingsInteractor,
+    private val communalSettingsInteractor: CommunalSettingsInteractor,
     private val appWidgetHost: CommunalAppWidgetHost,
     private val editWidgetsActivityStarter: EditWidgetsActivityStarter,
     private val userTracker: UserTracker,
@@ -358,7 +358,14 @@
     /** A list of widget content to be displayed in the communal hub. */
     val widgetContent: Flow<List<WidgetContent>> =
         combine(
-            widgetRepository.communalWidgets.map { filterWidgetsByExistingUsers(it) },
+            widgetRepository.communalWidgets
+                .map { filterWidgetsByExistingUsers(it) }
+                .combine(communalSettingsInteractor.allowedByDevicePolicyForWorkProfile) {
+                    // exclude widgets under work profile if not allowed by device policy
+                    widgets,
+                    allowedForWorkProfile ->
+                    filterWidgetsAllowedByDevicePolicy(widgets, allowedForWorkProfile)
+                },
             communalSettingsInteractor.communalWidgetCategories,
             updateOnWorkProfileBroadcastReceived,
         ) { widgets, allowedCategories, _ ->
@@ -380,6 +387,19 @@
             }
         }
 
+    /** Filter widgets based on whether their associated profile is allowed by device policy. */
+    private fun filterWidgetsAllowedByDevicePolicy(
+        list: List<CommunalWidgetContentModel>,
+        allowedByDevicePolicyForWorkProfile: Boolean
+    ): List<CommunalWidgetContentModel> =
+        if (allowedByDevicePolicyForWorkProfile) {
+            list
+        } else {
+            // Get associated work profile for the currently selected user.
+            val workProfile = userTracker.userProfiles.find { it.isManagedProfile }
+            list.filter { it.providerInfo.profile.identifier != workProfile?.id }
+        }
+
     /** A flow of available smartspace targets. Currently only showing timers. */
     private val smartspaceTargets: Flow<List<SmartspaceTarget>> =
         if (!smartspaceRepository.isSmartspaceRemoteViewsEnabled) {
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt
index 20f60b7..f9de609 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractor.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.communal.domain.interactor
 
+import android.content.pm.UserInfo
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.communal.data.model.CommunalEnabledState
 import com.android.systemui.communal.data.model.CommunalWidgetCategories
 import com.android.systemui.communal.data.repository.CommunalSettingsRepository
@@ -24,13 +26,18 @@
 import com.android.systemui.log.dagger.CommunalTableLog
 import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.log.table.logDiffsForTable
+import com.android.systemui.settings.UserTracker
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
+import java.util.concurrent.Executor
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 
@@ -40,8 +47,10 @@
 @Inject
 constructor(
     @Background private val bgScope: CoroutineScope,
+    @Background private val bgExecutor: Executor,
     private val repository: CommunalSettingsRepository,
     userInteractor: SelectedUserInteractor,
+    private val userTracker: UserTracker,
     @CommunalTableLog tableLogBuffer: TableLogBuffer,
 ) {
     /** Whether or not communal is enabled for the currently selected user. */
@@ -68,4 +77,33 @@
                 started = SharingStarted.Eagerly,
                 initialValue = CommunalWidgetCategories().categories
             )
+
+    private val workProfileUserInfoCallbackFlow: Flow<UserInfo?> = conflatedCallbackFlow {
+        fun send(profiles: List<UserInfo>) {
+            trySend(profiles.find { it.isManagedProfile })
+        }
+
+        val callback =
+            object : UserTracker.Callback {
+                override fun onProfilesChanged(profiles: List<UserInfo>) {
+                    send(profiles)
+                }
+            }
+        userTracker.addCallback(callback, bgExecutor)
+        send(userTracker.userProfiles)
+
+        awaitClose { userTracker.removeCallback(callback) }
+    }
+
+    /** Whether or not keyguard widgets are allowed for work profile by device policy manager. */
+    val allowedByDevicePolicyForWorkProfile: StateFlow<Boolean> =
+        workProfileUserInfoCallbackFlow
+            .flatMapLatest { workProfile ->
+                workProfile?.let { repository.getAllowedByDevicePolicy(it) } ?: flowOf(false)
+            }
+            .stateIn(
+                scope = bgScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = false
+            )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt
index b4773f6..cd2710e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalSettingsInteractorKosmos.kt
@@ -17,17 +17,21 @@
 package com.android.systemui.communal.domain.interactor
 
 import com.android.systemui.communal.data.repository.communalSettingsRepository
+import com.android.systemui.concurrency.fakeExecutor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.settings.userTracker
 import com.android.systemui.user.domain.interactor.selectedUserInteractor
 import com.android.systemui.util.mockito.mock
 
 val Kosmos.communalSettingsInteractor by Fixture {
     CommunalSettingsInteractor(
         bgScope = applicationCoroutineScope,
+        bgExecutor = fakeExecutor,
         repository = communalSettingsRepository,
         userInteractor = selectedUserInteractor,
+        userTracker = userTracker,
         tableLogBuffer = mock(),
     )
 }