diff --git a/res/layout/fingerprint_enroll_find_sensor_base.xml b/res/layout/fingerprint_enroll_find_sensor_base.xml
index da7fc30..62203f7 100644
--- a/res/layout/fingerprint_enroll_find_sensor_base.xml
+++ b/res/layout/fingerprint_enroll_find_sensor_base.xml
@@ -17,7 +17,6 @@
 
 <com.google.android.setupdesign.GlifLayout
     xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:app="http://schemas.android.com/apk/res-auto"
     android:id="@+id/setup_wizard_layout"
     style="?attr/fingerprint_layout_theme"
     android:layout_width="match_parent"
@@ -41,27 +40,5 @@
             android:layout_height="wrap_content"
             android:layout_gravity="center_horizontal"/>
 
-        <com.airbnb.lottie.LottieAnimationView
-            android:id="@+id/illustration_lottie"
-            android:layout_width="match_parent"
-            android:layout_height="wrap_content"
-            android:scaleType="centerInside"
-            android:visibility="gone"
-            app:lottie_imageAssetsFolder="images"
-            app:lottie_autoPlay="true"
-            app:lottie_loop="true"
-            app:lottie_rawRes="@raw/fingerprint_edu_lottie"/>
-
-        <com.airbnb.lottie.LottieAnimationView
-            android:id="@+id/illustration_lottie_portrait"
-            android:layout_width="match_parent"
-            android:layout_height="wrap_content"
-            android:scaleType="centerInside"
-            android:visibility="gone"
-            app:lottie_imageAssetsFolder="images"
-            app:lottie_autoPlay="true"
-            app:lottie_loop="true"
-            app:lottie_rawRes="@raw/fingerprint_edu_lottie_portrait"/>
-
     </LinearLayout>
 </com.google.android.setupdesign.GlifLayout>
diff --git a/res/raw/fingerprint_edu_lottie.json b/res/raw/fingerprint_edu_lottie.json
deleted file mode 100644
index e69de29..0000000
--- a/res/raw/fingerprint_edu_lottie.json
+++ /dev/null
diff --git a/res/raw/fingerprint_edu_lottie_portrait.json b/res/raw/fingerprint_edu_lottie_portrait.json
deleted file mode 100644
index e69de29..0000000
--- a/res/raw/fingerprint_edu_lottie_portrait.json
+++ /dev/null
diff --git a/src/com/android/settings/biometrics/fingerprint/FingerprintEnrollFindSensor.java b/src/com/android/settings/biometrics/fingerprint/FingerprintEnrollFindSensor.java
index 64e066e..3cea962 100644
--- a/src/com/android/settings/biometrics/fingerprint/FingerprintEnrollFindSensor.java
+++ b/src/com/android/settings/biometrics/fingerprint/FingerprintEnrollFindSensor.java
@@ -21,14 +21,13 @@
 import android.hardware.fingerprint.FingerprintManager;
 import android.hardware.fingerprint.FingerprintSensorPropertiesInternal;
 import android.os.Bundle;
-import android.view.OrientationEventListener;
-import android.view.Surface;
 import android.view.View;
 import android.view.View.OnClickListener;
 import android.view.accessibility.AccessibilityManager;
 
 import androidx.annotation.Nullable;
 
+import com.airbnb.lottie.LottieAnimationView;
 import com.android.settings.R;
 import com.android.settings.Utils;
 import com.android.settings.biometrics.BiometricEnrollBase;
@@ -36,7 +35,6 @@
 import com.android.settings.biometrics.BiometricUtils;
 import com.android.settings.password.ChooseLockSettingsHelper;
 
-import com.airbnb.lottie.LottieAnimationView;
 import com.google.android.setupcompat.template.FooterBarMixin;
 import com.google.android.setupcompat.template.FooterButton;
 
@@ -54,10 +52,6 @@
     private FingerprintEnrollSidecar mSidecar;
     private boolean mNextClicked;
     private boolean mCanAssumeUdfps;
-    private boolean mCanAssumeSidefps;
-
-    private OrientationEventListener mOrientationEventListener;
-    private int mPreviousRotation = 0;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -67,7 +61,6 @@
         final List<FingerprintSensorPropertiesInternal> props =
                 fingerprintManager.getSensorPropertiesInternal();
         mCanAssumeUdfps = props != null && props.size() == 1 && props.get(0).isAnyUdfpsType();
-        mCanAssumeSidefps = props != null && props.size() == 1 && props.get(0).isAnySidefpsType();
         setContentView(getContentView());
         mFooterBarMixin = getLayout().getMixin(FooterBarMixin.class);
         mFooterBarMixin.setSecondaryButton(
@@ -79,8 +72,6 @@
                         .build()
         );
 
-        listenOrientationEvent();
-
         if (mCanAssumeUdfps) {
             setHeaderText(R.string.security_settings_udfps_enroll_find_sensor_title);
             setDescriptionText(R.string.security_settings_udfps_enroll_find_sensor_message);
@@ -99,28 +90,6 @@
                 lottieAnimationView.setAnimation(R.raw.udfps_edu_a11y_lottie);
             }
 
-        } else if (mCanAssumeSidefps) {
-            setHeaderText(R.string.security_settings_fingerprint_enroll_find_sensor_title);
-            setDescriptionText(R.string.security_settings_fingerprint_enroll_find_sensor_message);
-            final LottieAnimationView lottieAnimationView = findViewById(R.id.illustration_lottie);
-            final LottieAnimationView lottieAnimationViewPortrait =
-                    findViewById(R.id.illustration_lottie_portrait);
-            final int rotation = getApplicationContext().getDisplay().getRotation();
-            switch(rotation) {
-                case Surface.ROTATION_90:
-                    lottieAnimationView.setVisibility(View.GONE);
-                    lottieAnimationViewPortrait.setVisibility(View.VISIBLE);
-                    break;
-                case Surface.ROTATION_270:
-                    lottieAnimationView.setVisibility(View.GONE);
-                    lottieAnimationViewPortrait.setVisibility(View.VISIBLE);
-                    lottieAnimationViewPortrait.setRotation(180);
-                    break;
-                default:
-                    lottieAnimationView.setVisibility(View.VISIBLE);
-                    lottieAnimationViewPortrait.setVisibility(View.GONE);
-                    break;
-            }
         } else {
             setHeaderText(R.string.security_settings_fingerprint_enroll_find_sensor_title);
             setDescriptionText(R.string.security_settings_fingerprint_enroll_find_sensor_message);
@@ -251,7 +220,6 @@
 
     @Override
     protected void onDestroy() {
-        stopListenOrientationEvent();
         super.onDestroy();
         if (mAnimation != null) {
             mAnimation.stopAnimation();
@@ -329,37 +297,4 @@
     public int getMetricsCategory() {
         return SettingsEnums.FINGERPRINT_FIND_SENSOR;
     }
-
-    private void listenOrientationEvent() {
-        if (!mCanAssumeSidefps) {
-            // Do nothing if the device doesn't support SideFPS.
-            return;
-        }
-        mOrientationEventListener = new OrientationEventListener(this) {
-            @Override
-            public void onOrientationChanged(int orientation) {
-                final int currentRotation = getDisplay().getRotation();
-                if ((mPreviousRotation == Surface.ROTATION_90
-                        && currentRotation == Surface.ROTATION_270) || (
-                        mPreviousRotation == Surface.ROTATION_270
-                                && currentRotation == Surface.ROTATION_90)) {
-                    mPreviousRotation = currentRotation;
-                    recreate();
-                }
-            }
-        };
-        mOrientationEventListener.enable();
-        mPreviousRotation = getDisplay().getRotation();
-    }
-
-    private void stopListenOrientationEvent() {
-        if (!mCanAssumeSidefps) {
-            // Do nothing if the device doesn't support SideFPS.
-            return;
-        }
-        if (mOrientationEventListener != null) {
-            mOrientationEventListener.disable();
-        }
-        mOrientationEventListener = null;
-    }
 }
diff --git a/src/com/android/settings/fuelgauge/AdvancedPowerUsageDetail.java b/src/com/android/settings/fuelgauge/AdvancedPowerUsageDetail.java
index de9f3cf..0dfa993 100644
--- a/src/com/android/settings/fuelgauge/AdvancedPowerUsageDetail.java
+++ b/src/com/android/settings/fuelgauge/AdvancedPowerUsageDetail.java
@@ -284,6 +284,15 @@
     }
 
     @Override
+    public void onPause() {
+        super.onPause();
+        if (mEnableTriState) {
+            Log.d(TAG, "Leave with mode: " + getSelectedPreference());
+            mBatteryOptimizeUtils.setAppUsageState(getSelectedPreference());
+        }
+    }
+
+    @Override
     public void onDestroy() {
         super.onDestroy();
         notifyBackupManager();
@@ -491,6 +500,18 @@
         mOptimizationMode = mBatteryOptimizeUtils.getAppOptimizationMode();
     }
 
+    private int getSelectedPreference() {
+        if (mRestrictedPreference.isChecked()) {
+            return BatteryOptimizeUtils.MODE_RESTRICTED;
+        } else if (mUnrestrictedPreference.isChecked()) {
+            return BatteryOptimizeUtils.MODE_UNRESTRICTED;
+        } else if (mOptimizePreference.isChecked()) {
+            return BatteryOptimizeUtils.MODE_OPTIMIZED;
+        } else {
+            return BatteryOptimizeUtils.MODE_UNKNOWN;
+        }
+    }
+
     private CharSequence getAppActiveTime(Bundle bundle) {
         final long foregroundTimeMs = bundle.getLong(EXTRA_FOREGROUND_TIME);
         final long backgroundTimeMs = bundle.getLong(EXTRA_BACKGROUND_TIME);
diff --git a/src/com/android/settings/fuelgauge/BatteryBackupHelper.java b/src/com/android/settings/fuelgauge/BatteryBackupHelper.java
index af16edb..fdbd12a 100644
--- a/src/com/android/settings/fuelgauge/BatteryBackupHelper.java
+++ b/src/com/android/settings/fuelgauge/BatteryBackupHelper.java
@@ -29,8 +29,8 @@
 import android.content.pm.UserInfo;
 import android.os.Build;
 import android.os.IDeviceIdleController;
-import android.os.RemoteException;
 import android.os.ParcelFileDescriptor;
+import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.UserHandle;
 import android.os.UserManager;
@@ -230,7 +230,7 @@
                 mBatteryOptimizeUtils != null
                         ? mBatteryOptimizeUtils /*testing only*/
                         : new BatteryOptimizeUtils(mContext, uid, packageName);
-        batteryOptimizeUtils.setAppOptimizationMode(mode);
+        batteryOptimizeUtils.setAppUsageState(mode);
         Log.d(TAG, String.format("restore:%s mode=%d", packageName, mode));
     }
 
diff --git a/src/com/android/settings/fuelgauge/BatteryOptimizeUtils.java b/src/com/android/settings/fuelgauge/BatteryOptimizeUtils.java
index 3d95395..9306a34 100644
--- a/src/com/android/settings/fuelgauge/BatteryOptimizeUtils.java
+++ b/src/com/android/settings/fuelgauge/BatteryOptimizeUtils.java
@@ -19,6 +19,7 @@
 import android.annotation.IntDef;
 import android.app.AppOpsManager;
 import android.content.Context;
+import android.os.AsyncTask;
 import android.util.Log;
 
 import androidx.annotation.VisibleForTesting;
@@ -89,36 +90,28 @@
         return getAppOptimizationMode(mMode, mAllowListed);
     }
 
-    public void setAppOptimizationMode(@OptimizationMode int mode) {
-        try {
-            setAppUsageStateInternal(mode);
-        } catch (Exception e) {
-            Log.e(TAG, "setAppUsageState() is failed for " + mPackageName, e);
-        }
-    }
-
     /** Sets the {@link OptimizationMode} for associated app. */
-    public void setAppUsageStateInternal(@OptimizationMode int mode) {
+    public void setAppUsageState(@OptimizationMode int mode) {
         if (getAppOptimizationMode(mMode, mAllowListed) == mode) {
             Log.w(TAG, "set the same optimization mode for: " + mPackageName);
             return;
         }
-        switch (mode) {
-            case MODE_RESTRICTED:
-                mBatteryUtils.setForceAppStandby(mUid, mPackageName, AppOpsManager.MODE_IGNORED);
-                mPowerAllowListBackend.removeApp(mPackageName);
-                break;
-            case MODE_UNRESTRICTED:
-                mBatteryUtils.setForceAppStandby(mUid, mPackageName, AppOpsManager.MODE_ALLOWED);
-                mPowerAllowListBackend.addApp(mPackageName);
-                break;
-            case MODE_OPTIMIZED:
-                mBatteryUtils.setForceAppStandby(mUid, mPackageName, AppOpsManager.MODE_ALLOWED);
-                mPowerAllowListBackend.removeApp(mPackageName);
-                break;
-            default:
-                Log.d(TAG, "set unknown app optimization mode.");
-        }
+
+        AsyncTask.execute(() -> {
+            switch (mode) {
+                case MODE_RESTRICTED:
+                    setAppOptimizationMode(AppOpsManager.MODE_IGNORED, /* allowListed */ false);
+                    break;
+                case MODE_UNRESTRICTED:
+                    setAppOptimizationMode(AppOpsManager.MODE_ALLOWED, /* allowListed */ true);
+                    break;
+                case MODE_OPTIMIZED:
+                    setAppOptimizationMode(AppOpsManager.MODE_ALLOWED, /* allowListed */ false);
+                    break;
+                default:
+                    Log.d(TAG, "set unknown app optimization mode.");
+            }
+        });
     }
 
     /**
@@ -142,6 +135,19 @@
         return mPackageName == null ? UNKNOWN_PACKAGE : mPackageName;
     }
 
+    private void setAppOptimizationMode(int appStandbyMode, boolean allowListed) {
+        try {
+            mBatteryUtils.setForceAppStandby(mUid, mPackageName, appStandbyMode);
+            if (allowListed) {
+                mPowerAllowListBackend.addApp(mPackageName);
+            } else {
+                mPowerAllowListBackend.removeApp(mPackageName);
+            }
+        } catch (Exception e) {
+            Log.e(TAG, "set OPTIMIZED failed for " + mPackageName, e);
+        }
+    }
+
     private void refreshState() {
         mPowerAllowListBackend.refreshList();
         mAllowListed = mPowerAllowListBackend.isAllowlisted(mPackageName);
diff --git a/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java b/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java
index d8bf27d..708ba5e 100644
--- a/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/OptimizedPreferenceController.java
@@ -73,12 +73,6 @@
 
     @Override
     public boolean handlePreferenceTreeClick(Preference preference) {
-        if (!KEY_OPTIMIZED_PREF.equals(preference.getKey())) {
-            return false;
-        }
-
-        mBatteryOptimizeUtils.setAppOptimizationMode(BatteryOptimizeUtils.MODE_OPTIMIZED);
-        Log.d(TAG, "Set optimized");
-        return true;
+        return getPreferenceKey().equals(preference.getKey());
     }
 }
diff --git a/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java b/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java
index ebe6563..99e742b 100644
--- a/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/RestrictedPreferenceController.java
@@ -76,12 +76,6 @@
 
     @Override
     public boolean handlePreferenceTreeClick(Preference preference) {
-        if (!KEY_RESTRICTED_PREF.equals(preference.getKey())) {
-            return false;
-        }
-
-        mBatteryOptimizeUtils.setAppOptimizationMode(BatteryOptimizeUtils.MODE_RESTRICTED);
-        Log.d(TAG, "Set restricted");
-        return true;
+        return getPreferenceKey().equals(preference.getKey());
     }
 }
diff --git a/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java b/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java
index f55be58..244a33c 100644
--- a/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/UnrestrictedPreferenceController.java
@@ -74,12 +74,6 @@
 
     @Override
     public boolean handlePreferenceTreeClick(Preference preference) {
-        if (!KEY_UNRESTRICTED_PREF.equals(preference.getKey())) {
-            return false;
-        }
-
-        mBatteryOptimizeUtils.setAppOptimizationMode(BatteryOptimizeUtils.MODE_UNRESTRICTED);
-        Log.d(TAG, "Set unrestricted");
-        return true;
+        return getPreferenceKey().equals(preference.getKey());
     }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/BatteryBackupHelperTest.java b/tests/robotests/src/com/android/settings/fuelgauge/BatteryBackupHelperTest.java
index 5ac8553..0f178e0 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/BatteryBackupHelperTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/BatteryBackupHelperTest.java
@@ -20,6 +20,7 @@
 import static com.android.settings.fuelgauge.BatteryBackupHelper.DELIMITER_MODE;
 import static com.android.settings.fuelgauge.BatteryOptimizeUtils.MODE_RESTRICTED;
 import static com.android.settings.fuelgauge.BatteryOptimizeUtils.MODE_UNRESTRICTED;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -33,7 +34,6 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
-import static org.mockito.Mockito.when;
 
 import android.app.AppOpsManager;
 import android.app.backup.BackupDataInputStream;
@@ -51,9 +51,6 @@
 
 import com.android.settingslib.fuelgauge.PowerAllowlistBackend;
 
-import java.util.Arrays;
-import java.util.List;
-
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -69,6 +66,10 @@
 import org.robolectric.annotation.Implements;
 import org.robolectric.annotation.Resetter;
 
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
 @RunWith(RobolectricTestRunner.class)
 @Config(shadows = {BatteryBackupHelperTest.ShadowUserHandle.class})
 public final class BatteryBackupHelperTest {
@@ -305,10 +306,11 @@
                 PACKAGE_NAME2 + DELIMITER_MODE + invalidNumberFormat;
 
         mBatteryBackupHelper.restoreOptimizationMode(packageModes.getBytes());
+        TimeUnit.SECONDS.sleep(1);
 
         final InOrder inOrder = inOrder(mBatteryOptimizeUtils);
-        inOrder.verify(mBatteryOptimizeUtils).setAppOptimizationMode(MODE_RESTRICTED);
-        inOrder.verify(mBatteryOptimizeUtils, never()).setAppOptimizationMode(anyInt());
+        inOrder.verify(mBatteryOptimizeUtils).setAppUsageState(MODE_RESTRICTED);
+        inOrder.verify(mBatteryOptimizeUtils, never()).setAppUsageState(anyInt());
     }
 
     @Test
@@ -319,11 +321,12 @@
                 PACKAGE_NAME3 + DELIMITER_MODE + MODE_RESTRICTED + DELIMITER;
 
         mBatteryBackupHelper.restoreOptimizationMode(packageModes.getBytes());
+        TimeUnit.SECONDS.sleep(1);
 
         final InOrder inOrder = inOrder(mBatteryOptimizeUtils);
-        inOrder.verify(mBatteryOptimizeUtils).setAppOptimizationMode(MODE_RESTRICTED);
-        inOrder.verify(mBatteryOptimizeUtils).setAppOptimizationMode(MODE_UNRESTRICTED);
-        inOrder.verify(mBatteryOptimizeUtils, never()).setAppOptimizationMode(MODE_RESTRICTED);
+        inOrder.verify(mBatteryOptimizeUtils).setAppUsageState(MODE_RESTRICTED);
+        inOrder.verify(mBatteryOptimizeUtils).setAppUsageState(MODE_UNRESTRICTED);
+        inOrder.verify(mBatteryOptimizeUtils, never()).setAppUsageState(MODE_RESTRICTED);
     }
 
     private void mockUid(int uid, String packageName) throws Exception {
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/BatteryOptimizeUtilsTest.java b/tests/robotests/src/com/android/settings/fuelgauge/BatteryOptimizeUtilsTest.java
index 53101c0..c154f42 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/BatteryOptimizeUtilsTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/BatteryOptimizeUtilsTest.java
@@ -42,6 +42,8 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 
+import java.util.concurrent.TimeUnit;
+
 @RunWith(RobolectricTestRunner.class)
 public class BatteryOptimizeUtilsTest {
 
@@ -125,12 +127,13 @@
     }
 
     @Test
-    public void testSetAppOptimizationMode_Restricted_verifyAction() {
+    public void testSetAppUsageState_Restricted_verifyAction() throws Exception {
         // Sets the current mode as MODE_UNRESTRICTED.
         mBatteryOptimizeUtils.mAllowListed = false;
         mBatteryOptimizeUtils.mMode = AppOpsManager.MODE_ALLOWED;
 
-        mBatteryOptimizeUtils.setAppOptimizationMode(MODE_RESTRICTED);
+        mBatteryOptimizeUtils.setAppUsageState(MODE_RESTRICTED);
+        TimeUnit.SECONDS.sleep(1);
 
         verify(mMockBatteryUtils).setForceAppStandby(UID,
                 PACKAGE_NAME, AppOpsManager.MODE_IGNORED);
@@ -138,8 +141,9 @@
     }
 
     @Test
-    public void testSetAppOptimizationMode_Unrestricted_verifyAction() {
-        mBatteryOptimizeUtils.setAppOptimizationMode(MODE_UNRESTRICTED);
+    public void testSetAppUsageState_Unrestricted_verifyAction() throws Exception {
+        mBatteryOptimizeUtils.setAppUsageState(MODE_UNRESTRICTED);
+        TimeUnit.SECONDS.sleep(1);
 
         verify(mMockBatteryUtils).setForceAppStandby(UID,
                 PACKAGE_NAME, AppOpsManager.MODE_ALLOWED);
@@ -147,8 +151,9 @@
     }
 
     @Test
-    public void testSetAppOptimizationMode_Optimized_verifyAction() {
-        mBatteryOptimizeUtils.setAppOptimizationMode(MODE_OPTIMIZED);
+    public void testSetAppUsageState_Optimized_verifyAction() throws Exception {
+        mBatteryOptimizeUtils.setAppUsageState(MODE_OPTIMIZED);
+        TimeUnit.SECONDS.sleep(1);
 
         verify(mMockBatteryUtils).setForceAppStandby(UID,
                 PACKAGE_NAME, AppOpsManager.MODE_ALLOWED);
@@ -156,12 +161,13 @@
     }
 
     @Test
-    public void testSetAppOptimizationMode_sameUnrestrictedMode_verifyNoAction() {
+    public void testSetAppUsageState_sameUnrestrictedMode_verifyNoAction() throws Exception {
         // Sets the current mode as MODE_UNRESTRICTED.
         mBatteryOptimizeUtils.mAllowListed = true;
         mBatteryOptimizeUtils.mMode = AppOpsManager.MODE_ALLOWED;
 
-        mBatteryOptimizeUtils.setAppOptimizationMode(MODE_UNRESTRICTED);
+        mBatteryOptimizeUtils.setAppUsageState(MODE_UNRESTRICTED);
+        TimeUnit.SECONDS.sleep(1);
 
         verifyZeroInteractions(mMockBackend);
         verifyZeroInteractions(mMockBatteryUtils);
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java
index 3f8aaeb..1b5ded0 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/OptimizedPreferenceControllerTest.java
@@ -18,8 +18,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import com.android.settingslib.widget.SelectorWithWidgetPreference;
@@ -98,14 +96,11 @@
         mPreference.setKey(mController.KEY_OPTIMIZED_PREF);
         mController.handlePreferenceTreeClick(mPreference);
 
-        verify(mockBatteryOptimizeUtils).setAppOptimizationMode(
-                BatteryOptimizeUtils.MODE_OPTIMIZED);
+        assertThat(mController.handlePreferenceTreeClick(mPreference)).isTrue();
     }
 
     @Test
     public void testHandlePreferenceTreeClick_incorrectPrefKey_noAction() {
-        mController.handlePreferenceTreeClick(mPreference);
-
-        verifyZeroInteractions(mockBatteryOptimizeUtils);
+        assertThat(mController.handlePreferenceTreeClick(mPreference)).isFalse();
     }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java
index 1776ad3..bacfb1a 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/RestrictedPreferenceControllerTest.java
@@ -18,8 +18,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import com.android.settingslib.widget.SelectorWithWidgetPreference;
@@ -106,14 +104,11 @@
         mPreference.setKey(mController.KEY_RESTRICTED_PREF);
         mController.handlePreferenceTreeClick(mPreference);
 
-        verify(mockBatteryOptimizeUtils).setAppOptimizationMode(
-                BatteryOptimizeUtils.MODE_RESTRICTED);
+        assertThat(mController.handlePreferenceTreeClick(mPreference)).isTrue();
     }
 
     @Test
     public void testHandlePreferenceTreeClick_incorrectPrefKey_noAction() {
-        mController.handlePreferenceTreeClick(mPreference);
-
-        verifyZeroInteractions(mockBatteryOptimizeUtils);
+        assertThat(mController.handlePreferenceTreeClick(mPreference)).isFalse();
     }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java
index 47d59b4..d5500fa 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/UnrestrictedPreferenceControllerTest.java
@@ -19,8 +19,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import com.android.settingslib.widget.SelectorWithWidgetPreference;
@@ -106,14 +104,11 @@
         mPreference.setKey(mController.KEY_UNRESTRICTED_PREF);
         mController.handlePreferenceTreeClick(mPreference);
 
-        verify(mockBatteryOptimizeUtils).setAppOptimizationMode(
-                BatteryOptimizeUtils.MODE_UNRESTRICTED);
+        assertThat(mController.handlePreferenceTreeClick(mPreference)).isTrue();
     }
 
     @Test
     public void testHandlePreferenceTreeClick_incorrectPrefKey_noAction() {
-        mController.handlePreferenceTreeClick(mPreference);
-
-        verifyZeroInteractions(mockBatteryOptimizeUtils);
+        assertThat(mController.handlePreferenceTreeClick(mPreference)).isFalse();
     }
 }