Add a flag to protect overwrite app optimization mode to restricted mode.

- Skip set app optimization mode to restricted if the flag disabled.

Bug: 348192536
Fix: 348192536
Test: atest
Flag: EXEMPT bug fix
Change-Id: I92d320ce1edfe9b7044afdcd058651b0d54bd6bd
diff --git a/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java b/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java
index 41ba519..e12cc9f 100644
--- a/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java
+++ b/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java
@@ -44,6 +44,9 @@
     /** Check whether the battery tips card is enabled in the battery usage page */
     boolean isBatteryTipsEnabled();
 
+    /** Check whether overwrite the app optimization mode to restricted mode is enabled */
+    boolean isRestrictedModeOverwriteEnabled();
+
     /** Check whether force expire the app optimization mode. */
     boolean isForceExpireAppOptimizationModeEnabled();
 
diff --git a/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java b/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
index 02dff7e..4f163a2 100644
--- a/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
+++ b/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
@@ -89,6 +89,11 @@
     }
 
     @Override
+    public boolean isRestrictedModeOverwriteEnabled() {
+        return false;
+    }
+
+    @Override
     public boolean isForceExpireAppOptimizationModeEnabled() {
         return false;
     }
diff --git a/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt b/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt
index afff1c9..508c4df 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt
+++ b/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt
@@ -76,8 +76,7 @@
     fun resetExpiredAppOptModeBeforeTimestamp(context: Context, queryTimestampMs: Long) =
         synchronized(appOptimizationModeLock) {
             val forceExpireEnabled =
-                featureFactory
-                    .powerUsageFeatureProvider.isForceExpireAppOptimizationModeEnabled
+                featureFactory.powerUsageFeatureProvider.isForceExpireAppOptimizationModeEnabled
             val eventsMap = getAppOptModeEventsMap(context)
             val expirationUids = ArrayList<Int>(eventsMap.size)
             for ((uid, event) in eventsMap) {
@@ -113,12 +112,22 @@
         getBatteryOptimizeUtils: (Int, String) -> BatteryOptimizeUtils,
     ) =
         synchronized(appOptimizationModeLock) {
+            val restrictedModeOverwriteEnabled =
+                featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled
             val eventsMap = getAppOptModeEventsMap(context)
             val expirationEvents: MutableMap<Int, AppOptimizationModeEvent> = ArrayMap()
             for (i in uids.indices) {
                 val uid = uids[i]
                 val packageName = packageNames[i]
                 val optimizationMode = optimizationModes[i]
+                if (
+                    !restrictedModeOverwriteEnabled &&
+                        optimizationMode == BatteryOptimizeUtils.MODE_RESTRICTED
+                ) {
+                    // Unable to set restricted mode due to flag protection.
+                    Log.w(TAG, "setOptimizationMode($packageName) into restricted ignored")
+                    continue
+                }
                 val originalOptMode: Int =
                     updateBatteryOptimizationMode(
                         context,
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java b/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
index cecf8f0..962c40f 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
@@ -78,6 +78,11 @@
     }
 
     @Test
+    public void isRestrictedModeOverwriteEnabled_returnFalse() {
+        assertThat(mPowerFeatureProvider.isRestrictedModeOverwriteEnabled()).isFalse();
+    }
+
+    @Test
     public void isForceExpireAppOptimizationModeEnabled_returnFalse() {
         assertThat(mPowerFeatureProvider.isForceExpireAppOptimizationModeEnabled()).isFalse();
     }
diff --git a/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt b/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt
index 94f082d..95f69da 100644
--- a/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt
+++ b/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt
@@ -83,8 +83,17 @@
     }
 
     @Test
-    fun updateAppOptModeExpirationInternal_withExpirationTime_verifyData() {
-        insertAppOptModeEventForTest(expirationTime = 1000L)
+    fun updateAppOptModeExpirationInternal_withoutExpirationTime_verifyEmptyList() {
+        insertAppOptModeEventForTest(expirationTime = UNLIMITED_EXPIRE_TIME)
+
+        assertThat(AppOptModeSharedPreferencesUtils.getAllEvents(context)).isEmpty()
+    }
+
+    @Test
+    fun updateAppOptModeExpirationInternal_setOptimizedModeWithFlagEnabled_verifyData() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(true)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_OPTIMIZED)
 
         val events = AppOptModeSharedPreferencesUtils.getAllEvents(context)
 
@@ -99,8 +108,46 @@
     }
 
     @Test
-    fun updateAppOptModeExpirationInternal_withoutExpirationTime_verifyEmptyList() {
-        insertAppOptModeEventForTest(expirationTime = UNLIMITED_EXPIRE_TIME)
+    fun updateAppOptModeExpirationInternal_setOptimizedModeWithFlagDisabled_verifyData() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(false)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_OPTIMIZED)
+
+        val events = AppOptModeSharedPreferencesUtils.getAllEvents(context)
+
+        assertThat(events).hasSize(1)
+        assertAppOptimizationModeEventInfo(
+            events[0],
+            UID,
+            PACKAGE_NAME,
+            MODE_OPTIMIZED,
+            expirationTime = 1000L
+        )
+    }
+
+    @Test
+    fun updateAppOptModeExpirationInternal_setRestrictedModeWithFlagEnabled_verifyData() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(true)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_RESTRICTED)
+
+        val events = AppOptModeSharedPreferencesUtils.getAllEvents(context)
+
+        assertThat(events).hasSize(1)
+        assertAppOptimizationModeEventInfo(
+            events[0],
+            UID,
+            PACKAGE_NAME,
+            MODE_RESTRICTED,
+            expirationTime = 1000L
+        )
+    }
+
+    @Test
+    fun updateAppOptModeExpirationInternal_setRestrictedModeWithFlagDisabled_verifyEmptyList() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(false)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_RESTRICTED)
 
         assertThat(AppOptModeSharedPreferencesUtils.getAllEvents(context)).isEmpty()
     }
@@ -237,14 +284,14 @@
         assertThat(currentOptMode).isEqualTo(MODE_RESTRICTED)
     }
 
-    private fun insertAppOptModeEventForTest(expirationTime: Long) {
+    private fun insertAppOptModeEventForTest(expirationTime: Long, mode: Int = MODE_OPTIMIZED) {
         whenever(testBatteryOptimizeUtils?.isOptimizeModeMutable).thenReturn(true)
-        whenever(testBatteryOptimizeUtils?.getAppOptimizationMode(true)).thenReturn(MODE_OPTIMIZED)
+        whenever(testBatteryOptimizeUtils?.getAppOptimizationMode(true)).thenReturn(mode)
         AppOptModeSharedPreferencesUtils.updateAppOptModeExpirationInternal(
             context,
             mutableListOf(UID),
             mutableListOf(PACKAGE_NAME),
-            mutableListOf(MODE_OPTIMIZED),
+            mutableListOf(mode),
             longArrayOf(expirationTime),
         ) { _: Int, _: String ->
             testBatteryOptimizeUtils