diff --git a/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java b/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java
index 708ba5efa51..88241b6ff39 100644
--- a/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java
@@ -53,16 +53,16 @@ public class OptimizedPreferenceController extends AbstractPreferenceController
             return;
         }
 
-        if (mBatteryOptimizeUtils.isSystemOrDefaultApp()) {
-            Log.d(TAG, "is system or default app, disable pref");
-            ((SelectorWithWidgetPreference) preference).setChecked(false);
-            preference.setEnabled(false);
-        } else if (mBatteryOptimizeUtils.getAppOptimizationMode()
+        if (mBatteryOptimizeUtils.getAppOptimizationMode()
                 == BatteryOptimizeUtils.MODE_OPTIMIZED) {
             Log.d(TAG, "is optimized states");
             ((SelectorWithWidgetPreference) preference).setChecked(true);
         } else {
             ((SelectorWithWidgetPreference) preference).setChecked(false);
+            if (mBatteryOptimizeUtils.isSystemOrDefaultApp()) {
+                Log.d(TAG, "is system or default app, disable pref");
+                preference.setEnabled(false);
+            }
         }
     }
 
diff --git a/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java b/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java
index 99e742ba236..fe896a6aed0 100644
--- a/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java
@@ -51,16 +51,16 @@ public class RestrictedPreferenceController extends AbstractPreferenceController
             preference.setEnabled(true);
         }
 
-        if (mBatteryOptimizeUtils.isSystemOrDefaultApp()) {
-            Log.d(TAG, "is system or default app, disable pref");
-            ((SelectorWithWidgetPreference) preference).setChecked(false);
-            preference.setEnabled(false);
-        } else if (mBatteryOptimizeUtils.getAppOptimizationMode()
+        if (mBatteryOptimizeUtils.getAppOptimizationMode()
                 == BatteryOptimizeUtils.MODE_RESTRICTED) {
             Log.d(TAG, "is restricted states");
             ((SelectorWithWidgetPreference) preference).setChecked(true);
         } else {
             ((SelectorWithWidgetPreference) preference).setChecked(false);
+            if (mBatteryOptimizeUtils.isSystemOrDefaultApp()) {
+                Log.d(TAG, "is system or default app, disable pref");
+                preference.setEnabled(false);
+            }
         }
     }
 
diff --git a/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java b/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java
index 244a33c4dfd..be4091c782f 100644
--- a/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java
@@ -50,15 +50,16 @@ public class UnrestrictedPreferenceController extends AbstractPreferenceControll
             preference.setEnabled(true);
         }
 
-        if (mBatteryOptimizeUtils.isSystemOrDefaultApp()) {
-            Log.d(TAG, "is system or default app, unrestricted states only");
-            ((SelectorWithWidgetPreference) preference).setChecked(true);
-        } else if (mBatteryOptimizeUtils.getAppOptimizationMode()
+        if (mBatteryOptimizeUtils.getAppOptimizationMode()
                 == BatteryOptimizeUtils.MODE_UNRESTRICTED) {
             Log.d(TAG, "is unrestricted states");
             ((SelectorWithWidgetPreference) preference).setChecked(true);
         } else {
             ((SelectorWithWidgetPreference) preference).setChecked(false);
+            if (mBatteryOptimizeUtils.isSystemOrDefaultApp()) {
+                Log.d(TAG, "is system or default app, disable pref");
+                preference.setEnabled(false);
+            }
         }
     }
 
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java
index 1b5ded0b7da..1fec92a8134 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java
@@ -60,6 +60,18 @@ public class OptimizedPreferenceControllerTest {
         assertThat(mPreference.isChecked()).isTrue();
     }
 
+    @Test
+    public void testUpdateState_isSystemOrDefaultAppAndOptimizeStates_prefChecked() {
+        when(mockBatteryOptimizeUtils.isValidPackageName()).thenReturn(true);
+        when(mockBatteryOptimizeUtils.isSystemOrDefaultApp()).thenReturn(true);
+        when(mockBatteryOptimizeUtils.getAppOptimizationMode()).thenReturn(
+                BatteryOptimizeUtils.MODE_OPTIMIZED);
+
+        mController.updateState(mPreference);
+
+        assertThat(mPreference.isChecked()).isTrue();
+    }
+
     @Test
     public void testUpdateState_isSystemOrDefaultApp_prefUnchecked() {
         when(mockBatteryOptimizeUtils.isValidPackageName()).thenReturn(true);
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java
index bacfb1af714..944376c77dd 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java
@@ -69,7 +69,19 @@ public class RestrictedPreferenceControllerTest {
     }
 
     @Test
-    public void testUpdateState_isSystemOrDefaultApp_prefChecked() {
+    public void testUpdateState_isSystemOrDefaultAppAndRestrictedStates_prefChecked() {
+        when(mockBatteryOptimizeUtils.isValidPackageName()).thenReturn(true);
+        when(mockBatteryOptimizeUtils.isSystemOrDefaultApp()).thenReturn(true);
+        when(mockBatteryOptimizeUtils.getAppOptimizationMode()).thenReturn(
+                BatteryOptimizeUtils.MODE_RESTRICTED);
+
+        mController.updateState(mPreference);
+
+        assertThat(mPreference.isChecked()).isTrue();
+    }
+
+    @Test
+    public void testUpdateState_isSystemOrDefaultApp_prefUnchecked() {
         when(mockBatteryOptimizeUtils.isValidPackageName()).thenReturn(true);
         when(mockBatteryOptimizeUtils.isSystemOrDefaultApp()).thenReturn(true);
 
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java
index d5500faf983..c5642df576c 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java
@@ -70,13 +70,26 @@ public class UnrestrictedPreferenceControllerTest {
     }
 
     @Test
-    public void testUpdateState_isSystemOrDefaultApp_prefChecked() {
+    public void testUpdateState_isSystemOrDefaultAppAndUnrestrictedStates_prefChecked() {
+        when(mockBatteryOptimizeUtils.isValidPackageName()).thenReturn(true);
+        when(mockBatteryOptimizeUtils.isSystemOrDefaultApp()).thenReturn(true);
+        when(mockBatteryOptimizeUtils.getAppOptimizationMode()).thenReturn(
+                BatteryOptimizeUtils.MODE_UNRESTRICTED);
+
+        mController.updateState(mPreference);
+
+        assertThat(mPreference.isChecked()).isTrue();
+    }
+
+    @Test
+    public void testUpdateState_isSystemOrDefaultApp_prefUnchecked() {
         when(mockBatteryOptimizeUtils.isValidPackageName()).thenReturn(true);
         when(mockBatteryOptimizeUtils.isSystemOrDefaultApp()).thenReturn(true);
 
         mController.updateState(mPreference);
 
-        assertThat(mPreference.isChecked()).isTrue();
+        assertThat(mPreference.isChecked()).isFalse();
+        assertThat(mPreference.isEnabled()).isFalse();
     }
 
     @Test