Merge "Add base user restriction enforcement to new brightness" into main

4 files changed, 121 insertions, 10 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepositoryImplTest.kt
index 312c14df9505..fec56ed919de 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepositoryImplTest.kt
@@ -18,9 +18,13 @@ package com.android.systemui.brightness.data.repository
 
 import android.content.applicationContext
 import android.os.UserManager
-import androidx.test.ext.junit.runners.AndroidJUnit4
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.FlagsParameterization
+import android.platform.test.flag.junit.FlagsParameterization.allCombinationsOf
 import androidx.test.filters.SmallTest
 import com.android.settingslib.RestrictedLockUtils.EnforcedAdmin
+import com.android.systemui.Flags.FLAG_ENFORCE_BRIGHTNESS_BASE_USER_RESTRICTION
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testDispatcher
@@ -40,14 +44,18 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
 
 @SmallTest
-@RunWith(AndroidJUnit4::class)
-class BrightnessPolicyRepositoryImplTest : SysuiTestCase() {
+@RunWith(ParameterizedAndroidJunit4::class)
+class BrightnessPolicyRepositoryImplTest(flags: FlagsParameterization) : SysuiTestCase() {
 
-    private val kosmos = testKosmos()
+    init {
+        mSetFlagsRule.setFlagsParameterization(flags)
+    }
 
-    private val fakeUserRepository = kosmos.fakeUserRepository
+    private val kosmos = testKosmos()
 
     private val mockUserRestrictionChecker: UserRestrictionChecker = mock {
         whenever(checkIfRestrictionEnforced(any(), anyString(), anyInt())).thenReturn(null)
@@ -130,7 +138,83 @@ class BrightnessPolicyRepositoryImplTest : SysuiTestCase() {
             }
         }
 
-    private companion object {
-        val RESTRICTION = UserManager.DISALLOW_CONFIG_BRIGHTNESS
+    @Test
+    @DisableFlags(FLAG_ENFORCE_BRIGHTNESS_BASE_USER_RESTRICTION)
+    fun brightnessBaseUserRestriction_flagOff_noRestriction() =
+        with(kosmos) {
+            testScope.runTest {
+                whenever(
+                        mockUserRestrictionChecker.hasBaseUserRestriction(
+                            any(),
+                            eq(RESTRICTION),
+                            eq(userRepository.getSelectedUserInfo().id)
+                        )
+                    )
+                    .thenReturn(true)
+
+                val restrictions by collectLastValue(underTest.restrictionPolicy)
+
+                assertThat(restrictions).isEqualTo(PolicyRestriction.NoRestriction)
+            }
+        }
+
+    @Test
+    fun bothRestrictions_returnsSetEnforcedAdminFromCheck() =
+        with(kosmos) {
+            testScope.runTest {
+                val enforcedAdmin: EnforcedAdmin =
+                    EnforcedAdmin.createDefaultEnforcedAdminWithRestriction(RESTRICTION)
+
+                whenever(
+                        mockUserRestrictionChecker.checkIfRestrictionEnforced(
+                            any(),
+                            eq(RESTRICTION),
+                            eq(userRepository.getSelectedUserInfo().id)
+                        )
+                    )
+                    .thenReturn(enforcedAdmin)
+
+                whenever(
+                        mockUserRestrictionChecker.hasBaseUserRestriction(
+                            any(),
+                            eq(RESTRICTION),
+                            eq(userRepository.getSelectedUserInfo().id)
+                        )
+                    )
+                    .thenReturn(true)
+
+                val restrictions by collectLastValue(underTest.restrictionPolicy)
+
+                assertThat(restrictions).isEqualTo(PolicyRestriction.Restricted(enforcedAdmin))
+            }
+        }
+
+    @Test
+    @EnableFlags(FLAG_ENFORCE_BRIGHTNESS_BASE_USER_RESTRICTION)
+    fun brightnessBaseUserRestriction_flagOn_emptyRestriction() =
+        with(kosmos) {
+            testScope.runTest {
+                whenever(
+                        mockUserRestrictionChecker.hasBaseUserRestriction(
+                            any(),
+                            eq(RESTRICTION),
+                            eq(userRepository.getSelectedUserInfo().id)
+                        )
+                    )
+                    .thenReturn(true)
+
+                val restrictions by collectLastValue(underTest.restrictionPolicy)
+
+                assertThat(restrictions).isEqualTo(PolicyRestriction.Restricted(EnforcedAdmin()))
+            }
+        }
+
+    companion object {
+        private const val RESTRICTION = UserManager.DISALLOW_CONFIG_BRIGHTNESS
+        @JvmStatic
+        @Parameters(name = "{0}")
+        fun getParams(): List<FlagsParameterization> {
+            return allCombinationsOf(FLAG_ENFORCE_BRIGHTNESS_BASE_USER_RESTRICTION)
+        }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/domain/interactor/BrightnessPolicyEnforcementInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/domain/interactor/BrightnessPolicyEnforcementInteractorTest.kt
index 85a4bcf62223..11f523846268 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/domain/interactor/BrightnessPolicyEnforcementInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/brightness/domain/interactor/BrightnessPolicyEnforcementInteractorTest.kt
@@ -48,7 +48,6 @@ class BrightnessPolicyEnforcementInteractorTest : SysuiTestCase() {
     private val kosmos = testKosmos()
 
     private val mockActivityStarter = kosmos.activityStarter
-    private val fakeBrightnessPolicyEnforcementInteractor = kosmos.fakeBrightnessPolicyRepository
 
     private val underTest =
         with(kosmos) {
@@ -70,7 +69,18 @@ class BrightnessPolicyEnforcementInteractorTest : SysuiTestCase() {
 
                 fakeBrightnessPolicyRepository.setCurrentUserRestricted()
 
-                assertThat(restriction).isInstanceOf(PolicyRestriction.Restricted::class.java)
+                assertThat(restriction)
+                    .isEqualTo(
+                        PolicyRestriction.Restricted(
+                            EnforcedAdmin.createDefaultEnforcedAdminWithRestriction(
+                                BrightnessPolicyRepository.RESTRICTION
+                            )
+                        )
+                    )
+
+                fakeBrightnessPolicyRepository.setBaseUserRestriction()
+
+                assertThat(restriction).isEqualTo(PolicyRestriction.Restricted(EnforcedAdmin()))
             }
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepository.kt b/packages/SystemUI/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepository.kt
index c018ecb25835..0544a4f66295 100644
--- a/packages/SystemUI/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/brightness/data/repository/BrightnessPolicyRepository.kt
@@ -18,6 +18,8 @@ package com.android.systemui.brightness.data.repository
 
 import android.content.Context
 import android.os.UserManager
+import com.android.settingslib.RestrictedLockUtils
+import com.android.systemui.Flags.enforceBrightnessBaseUserRestriction
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
@@ -66,7 +68,18 @@ constructor(
                         user.id
                     )
                     ?.let { PolicyRestriction.Restricted(it) }
-                    ?: PolicyRestriction.NoRestriction
+                    ?: if (
+                        enforceBrightnessBaseUserRestriction() &&
+                            userRestrictionChecker.hasBaseUserRestriction(
+                                applicationContext,
+                                UserManager.DISALLOW_CONFIG_BRIGHTNESS,
+                                user.id
+                            )
+                    ) {
+                        PolicyRestriction.Restricted(RestrictedLockUtils.EnforcedAdmin())
+                    } else {
+                        PolicyRestriction.NoRestriction
+                    }
             }
             .flowOn(backgroundDispatcher)
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/brightness/data/repository/FakeBrightnessPolicyRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/brightness/data/repository/FakeBrightnessPolicyRepository.kt
index d3ceb15ca5a2..f5d02f328336 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/brightness/data/repository/FakeBrightnessPolicyRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/brightness/data/repository/FakeBrightnessPolicyRepository.kt
@@ -38,4 +38,8 @@ class FakeBrightnessPolicyRepository : BrightnessPolicyRepository {
                 )
             )
     }
+
+    fun setBaseUserRestriction() {
+        _restrictionPolicy.value = PolicyRestriction.Restricted(RestrictedLockUtils.EnforcedAdmin())
+    }
 }