From 92757b2454c3bbc6d54a5b49efad4ebafc9fe01b Mon Sep 17 00:00:00 2001
From: jackqdyulei <jackqdyulei@google.com>
Date: Wed, 24 Jan 2018 19:42:42 -0800
Subject: [PATCH] Add controller for battery saver button

Bug: 72228477
Test: RunSettingsRoboTests
Change-Id: Iae9096a10553ea761d099e1e72c1de56b693bac0
---
 ...ttings_button.xml => two_state_button.xml} |   6 +-
 res/values/attrs.xml                          |   8 ++
 res/values/styles_preference.xml              |   4 +
 res/values/themes.xml                         |   1 +
 res/xml/battery_saver_settings.xml            |   5 +-
 .../applications/LayoutPreference.java        |  31 +++--
 .../fuelgauge/BatteryBroadcastReceiver.java   |  11 +-
 ...atterySaverButtonPreferenceController.java |  85 ++++++++++++
 .../batterysaver/BatterySaverSettings.java    |   1 +
 .../widget/TwoStateButtonPreference.java      |  62 +++++++++
 .../TwoStateButtonPreferenceController.java   |  73 +++++++++++
 ...rySaverButtonPreferenceControllerTest.java | 121 ++++++++++++++++++
 .../testutils/shadow/ShadowPowerManager.java  |  32 +++++
 ...woStateButtonPreferenceControllerTest.java | 105 +++++++++++++++
 14 files changed, 526 insertions(+), 19 deletions(-)
 rename res/layout/{battery_saver_settings_button.xml => two_state_button.xml} (85%)
 create mode 100644 src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceController.java
 create mode 100644 src/com/android/settings/widget/TwoStateButtonPreference.java
 create mode 100644 src/com/android/settings/widget/TwoStateButtonPreferenceController.java
 create mode 100644 tests/robotests/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceControllerTest.java
 create mode 100644 tests/robotests/src/com/android/settings/testutils/shadow/ShadowPowerManager.java
 create mode 100644 tests/robotests/src/com/android/settings/widget/TwoStateButtonPreferenceControllerTest.java

diff --git a/res/layout/battery_saver_settings_button.xml b/res/layout/two_state_button.xml
similarity index 85%
rename from res/layout/battery_saver_settings_button.xml
rename to res/layout/two_state_button.xml
index 015d4d88f86..22bf2bd3da0 100644
--- a/res/layout/battery_saver_settings_button.xml
+++ b/res/layout/two_state_button.xml
@@ -25,23 +25,21 @@
     android:layout_height="match_parent">
 
     <Button
-        android:id="@+id/battery_saver_on_button"
+        android:id="@+id/state_on_button"
         style="@style/ActionPrimaryButton"
         android:layout_width="0dp"
         android:layout_weight="1"
         android:layout_height="wrap_content"
         android:layout_gravity="center"
-        android:text="@string/battery_saver_button_turn_on"
         android:paddingEnd="8dp" />
 
     <Button
-        android:id="@+id/battery_saver_off_button"
+        android:id="@+id/state_off_button"
         style="@style/ActionSecondaryButton"
         android:layout_width="0dp"
         android:layout_weight="1"
         android:layout_height="wrap_content"
         android:layout_gravity="center"
-        android:text="@string/battery_saver_button_turn_off"
         android:paddingEnd="8dp" />
 
 </LinearLayout>
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 2d48e423e68..d37091c5948 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -162,4 +162,12 @@
         <attr name="showPercentString" format="boolean" />
         <attr name="thickness" format="dimension" />
     </declare-styleable>
+
+    <!-- For TwoStatesButtonPreference -->
+    <declare-styleable name="TwoStateButtonPreference">
+        <attr name="textOn" format="reference" />
+        <attr name="textOff" format="reference" />
+    </declare-styleable>
+
+    <attr name="twoStateButtonPreferenceStyle" format="reference" />
 </resources>
diff --git a/res/values/styles_preference.xml b/res/values/styles_preference.xml
index 056d55a0997..5fe1ef41787 100644
--- a/res/values/styles_preference.xml
+++ b/res/values/styles_preference.xml
@@ -36,4 +36,8 @@
         <item name="android:widgetLayout">@layout/preference_widget_sync_toggle</item>
     </style>
 
+    <style name="TwoStateButtonPreference" parent="Preference.SettingsBase">
+        <item name="android:layout">@layout/two_state_button</item>
+    </style>
+
 </resources>
\ No newline at end of file
diff --git a/res/values/themes.xml b/res/values/themes.xml
index d960654f11c..a8ceca3dbf6 100644
--- a/res/values/themes.xml
+++ b/res/values/themes.xml
@@ -176,6 +176,7 @@
         <item name="preferenceFragmentStyle">@style/SettingsPreferenceFragmentStyle</item>
         <item name="apnPreferenceStyle">@style/ApnPreference</item>
         <item name="seekBarPreferenceStyle">@style/SettingsSeekBarPreference</item>
+        <item name="twoStateButtonPreferenceStyle">@style/TwoStateButtonPreference</item>
     </style>
 
     <style name="PreferenceTheme.SetupWizard" parent="PreferenceTheme">
diff --git a/res/xml/battery_saver_settings.xml b/res/xml/battery_saver_settings.xml
index 9e32440cba0..43971991b6f 100644
--- a/res/xml/battery_saver_settings.xml
+++ b/res/xml/battery_saver_settings.xml
@@ -33,10 +33,11 @@
         android:max="75"
         android:min="5"/>
 
-    <com.android.settings.applications.LayoutPreference
+    <com.android.settings.widget.TwoStateButtonPreference
         android:key="battery_saver_button_container"
         android:selectable="false"
-        android:layout="@layout/battery_saver_settings_button"/>
+        settings:textOn="@string/battery_saver_button_turn_on"
+        settings:textOff="@string/battery_saver_button_turn_off"/>
 
     <PreferenceCategory
         android:key="battery_saver_footer">
diff --git a/src/com/android/settings/applications/LayoutPreference.java b/src/com/android/settings/applications/LayoutPreference.java
index f2bd183bfb0..6940a4bd87d 100644
--- a/src/com/android/settings/applications/LayoutPreference.java
+++ b/src/com/android/settings/applications/LayoutPreference.java
@@ -42,6 +42,24 @@ public class LayoutPreference extends Preference {
 
     public LayoutPreference(Context context, AttributeSet attrs) {
         super(context, attrs);
+        init(context, attrs, 0 /* defStyleAttr */);
+    }
+
+    public LayoutPreference(Context context, AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+        init(context, attrs, defStyleAttr);
+    }
+
+    public LayoutPreference(Context context, int resource) {
+        this(context, LayoutInflater.from(context).inflate(resource, null, false));
+    }
+
+    public LayoutPreference(Context context, View view) {
+        super(context);
+        setView(view);
+    }
+
+    private void init(Context context, AttributeSet attrs, int defStyleAttr) {
         TypedArray a = context.obtainStyledAttributes(attrs, R.styleable.Preference);
         mAllowDividerAbove = TypedArrayUtils.getBoolean(a, R.styleable.Preference_allowDividerAbove,
                 R.styleable.Preference_allowDividerAbove, false);
@@ -50,7 +68,7 @@ public class LayoutPreference extends Preference {
         a.recycle();
 
         a = context.obtainStyledAttributes(
-                attrs, com.android.internal.R.styleable.Preference, 0, 0);
+                attrs, com.android.internal.R.styleable.Preference, defStyleAttr, 0);
         int layoutResource = a.getResourceId(com.android.internal.R.styleable.Preference_layout,
                 0);
         if (layoutResource == 0) {
@@ -64,15 +82,6 @@ public class LayoutPreference extends Preference {
         setView(view);
     }
 
-    public LayoutPreference(Context context, int resource) {
-        this(context, LayoutInflater.from(context).inflate(resource, null, false));
-    }
-
-    public LayoutPreference(Context context, View view) {
-        super(context);
-        setView(view);
-    }
-
     private void setView(View view) {
         setLayoutResource(R.layout.layout_preference_frame);
         final ViewGroup allDetails = view.findViewById(R.id.all_details);
@@ -106,4 +115,4 @@ public class LayoutPreference extends Preference {
         return mRootView.findViewById(id);
     }
 
-}
+}
\ No newline at end of file
diff --git a/src/com/android/settings/fuelgauge/BatteryBroadcastReceiver.java b/src/com/android/settings/fuelgauge/BatteryBroadcastReceiver.java
index d0f4080a4e2..ca4742d1fcb 100644
--- a/src/com/android/settings/fuelgauge/BatteryBroadcastReceiver.java
+++ b/src/com/android/settings/fuelgauge/BatteryBroadcastReceiver.java
@@ -27,7 +27,7 @@ import com.android.settings.Utils;
 
 /**
  * Use this broadcastReceiver to listen to the battery change, and it will invoke
- * {@link OnBatteryChangedListener} if any of the followings has been changed:
+ * {@link OnBatteryChangedListener} if any of the following has been changed:
  *
  * 1. Battery level(e.g. 100%->99%)
  * 2. Battery status(e.g. plugged->unplugged)
@@ -35,7 +35,14 @@ import com.android.settings.Utils;
  */
 public class BatteryBroadcastReceiver extends BroadcastReceiver {
 
-    interface OnBatteryChangedListener {
+    /**
+     * Callback when the following has been changed:
+     *
+     * Battery level(e.g. 100%->99%)
+     * Battery status(e.g. plugged->unplugged)
+     * Battery saver(e.g. off->on)
+     */
+    public interface OnBatteryChangedListener {
         void onBatteryChanged();
     }
 
diff --git a/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceController.java b/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceController.java
new file mode 100644
index 00000000000..17b42b80f69
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceController.java
@@ -0,0 +1,85 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.fuelgauge.batterysaver;
+
+import android.content.Context;
+import android.os.PowerManager;
+import android.support.annotation.VisibleForTesting;
+import android.support.v7.preference.Preference;
+
+import com.android.settings.fuelgauge.BatteryBroadcastReceiver;
+import com.android.settings.widget.TwoStateButtonPreferenceController;
+import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.core.lifecycle.LifecycleObserver;
+import com.android.settingslib.core.lifecycle.events.OnStart;
+import com.android.settingslib.core.lifecycle.events.OnStop;
+
+/**
+ * Controller to update the battery saver button
+ */
+//TODO(b/72228477): disable the button if device is charging.
+public class BatterySaverButtonPreferenceController extends
+        TwoStateButtonPreferenceController implements
+        LifecycleObserver, OnStart, OnStop, BatteryBroadcastReceiver.OnBatteryChangedListener {
+    private static final String KEY = "battery_saver_button_container";
+    private BatteryBroadcastReceiver mBatteryBroadcastReceiver;
+    @VisibleForTesting
+    PowerManager mPowerManager;
+
+    public BatterySaverButtonPreferenceController(Context context, Lifecycle lifecycle) {
+        super(context, KEY);
+        mPowerManager = (PowerManager) context.getSystemService(Context.POWER_SERVICE);
+        mBatteryBroadcastReceiver = new BatteryBroadcastReceiver(context);
+        mBatteryBroadcastReceiver.setBatteryChangedListener(this);
+        if (lifecycle != null) {
+            lifecycle.addObserver(this);
+        }
+    }
+
+    @Override
+    public void onStart() {
+        mBatteryBroadcastReceiver.register();
+    }
+
+    @Override
+    public void onStop() {
+        mBatteryBroadcastReceiver.unRegister();
+    }
+
+    @Override
+    public void updateState(Preference preference) {
+        super.updateState(preference);
+        final boolean lowPowerModeOn = mPowerManager.isPowerSaveMode();
+        updateButton(!lowPowerModeOn);
+    }
+
+    @Override
+    public int getAvailabilityStatus() {
+        return AVAILABLE;
+    }
+
+    @Override
+    public void onBatteryChanged() {
+        final boolean lowPowerModeOn = mPowerManager.isPowerSaveMode();
+        updateButton(!lowPowerModeOn);
+    }
+
+    @Override
+    public void onButtonClicked(boolean stateOn) {
+        mPowerManager.setPowerSaveMode(stateOn);
+    }
+}
\ No newline at end of file
diff --git a/src/com/android/settings/fuelgauge/batterysaver/BatterySaverSettings.java b/src/com/android/settings/fuelgauge/batterysaver/BatterySaverSettings.java
index 392032ce21b..2340675554d 100644
--- a/src/com/android/settings/fuelgauge/batterysaver/BatterySaverSettings.java
+++ b/src/com/android/settings/fuelgauge/batterysaver/BatterySaverSettings.java
@@ -78,6 +78,7 @@ public class BatterySaverSettings extends DashboardFragment {
         final List<AbstractPreferenceController> controllers = new ArrayList<>();
         controllers.add(new AutoBatterySaverPreferenceController(context));
         controllers.add(new AutoBatterySeekBarPreferenceController(context, lifecycle));
+        controllers.add(new BatterySaverButtonPreferenceController(context, lifecycle));
         return controllers;
     }
 
diff --git a/src/com/android/settings/widget/TwoStateButtonPreference.java b/src/com/android/settings/widget/TwoStateButtonPreference.java
new file mode 100644
index 00000000000..6b5fbbb0e11
--- /dev/null
+++ b/src/com/android/settings/widget/TwoStateButtonPreference.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.widget;
+
+import android.content.Context;
+import android.content.res.TypedArray;
+import android.support.v4.content.res.TypedArrayUtils;
+import android.util.AttributeSet;
+import android.widget.Button;
+
+import com.android.settings.R;
+import com.android.settings.applications.LayoutPreference;
+
+/**
+ * Preference that presents a button with two states(On vs Off)
+ */
+public class TwoStateButtonPreference extends LayoutPreference {
+    public TwoStateButtonPreference(Context context, AttributeSet attrs) {
+        super(context, attrs, TypedArrayUtils.getAttr(
+                context, R.attr.twoStateButtonPreferenceStyle, android.R.attr.preferenceStyle));
+
+        if (attrs != null) {
+            final TypedArray styledAttrs = context.obtainStyledAttributes(attrs,
+                    R.styleable.TwoStateButtonPreference);
+            final int textOnId = styledAttrs.getResourceId(
+                    R.styleable.TwoStateButtonPreference_textOn,
+                    R.string.summary_placeholder);
+            final int textOffId = styledAttrs.getResourceId(
+                    R.styleable.TwoStateButtonPreference_textOff,
+                    R.string.summary_placeholder);
+            styledAttrs.recycle();
+
+            final Button buttonOn = getStateOnButton();
+            buttonOn.setText(textOnId);
+            final Button buttonOff = getStateOffButton();
+            buttonOff.setText(textOffId);
+        }
+    }
+
+    public Button getStateOnButton() {
+        return findViewById(R.id.state_on_button);
+    }
+
+
+    public Button getStateOffButton() {
+        return findViewById(R.id.state_off_button);
+    }
+}
\ No newline at end of file
diff --git a/src/com/android/settings/widget/TwoStateButtonPreferenceController.java b/src/com/android/settings/widget/TwoStateButtonPreferenceController.java
new file mode 100644
index 00000000000..5f5e0984699
--- /dev/null
+++ b/src/com/android/settings/widget/TwoStateButtonPreferenceController.java
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.widget;
+
+import android.content.Context;
+import android.support.v7.preference.PreferenceScreen;
+import android.view.View;
+import android.widget.Button;
+
+import com.android.settings.R;
+import com.android.settings.core.BasePreferenceController;
+
+/**
+ * Controller to update the button with two states(On vs Off).
+ */
+public abstract class TwoStateButtonPreferenceController extends BasePreferenceController
+        implements View.OnClickListener {
+    private Button mButtonOn;
+    private Button mButtonOff;
+
+    public TwoStateButtonPreferenceController(Context context, String key) {
+        super(context, key);
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        final TwoStateButtonPreference preference =
+                (TwoStateButtonPreference) screen.findPreference(getPreferenceKey());
+        mButtonOn = preference.getStateOnButton();
+        mButtonOn.setOnClickListener(this);
+        mButtonOff = preference.getStateOffButton();
+        mButtonOff.setOnClickListener(this);
+    }
+
+    protected void updateButton(boolean stateOn) {
+        if (stateOn) {
+            mButtonOff.setVisibility(View.GONE);
+            mButtonOn.setVisibility(View.VISIBLE);
+        } else {
+            mButtonOff.setVisibility(View.VISIBLE);
+            mButtonOn.setVisibility(View.GONE);
+        }
+    }
+
+    @Override
+    public void onClick(View v) {
+        final boolean stateOn = v.getId() == R.id.state_on_button;
+        onButtonClicked(stateOn);
+    }
+
+    /**
+     * Callback when button is clicked
+     *
+     * @param stateOn {@code true} if stateOn button is clicked, otherwise it means stateOff
+     *                button is clicked
+     */
+    public abstract void onButtonClicked(boolean stateOn);
+}
\ No newline at end of file
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceControllerTest.java
new file mode 100644
index 00000000000..29f924f0786
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batterysaver/BatterySaverButtonPreferenceControllerTest.java
@@ -0,0 +1,121 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.settings.fuelgauge.batterysaver;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+import android.arch.lifecycle.LifecycleOwner;
+import android.content.Context;
+import android.os.PowerManager;
+import android.support.v7.preference.PreferenceScreen;
+import android.view.View;
+import android.widget.Button;
+
+import com.android.settings.R;
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+import com.android.settings.widget.TwoStateButtonPreference;
+import com.android.settingslib.core.lifecycle.Lifecycle;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.Shadows;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowPowerManager;
+
+@RunWith(SettingsRobolectricTestRunner.class)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION,
+        shadows = com.android.settings.testutils.shadow.ShadowPowerManager.class)
+public class BatterySaverButtonPreferenceControllerTest {
+    private BatterySaverButtonPreferenceController mController;
+    private Context mContext;
+    private Lifecycle mLifecycle;
+    private LifecycleOwner mLifecycleOwner;
+    private Button mButtonOn;
+    private Button mButtonOff;
+    private ShadowPowerManager mShadowPowerManager;
+    @Mock
+    private TwoStateButtonPreference mPreference;
+    @Mock
+    private PreferenceScreen mPreferenceScreen;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mLifecycleOwner = () -> mLifecycle;
+        mLifecycle = new Lifecycle(mLifecycleOwner);
+        mContext = spy(RuntimeEnvironment.application);
+        PowerManager powerManager = (PowerManager) mContext.getSystemService(Context.POWER_SERVICE);
+        mShadowPowerManager = Shadows.shadowOf(powerManager);
+        doReturn(mPreference).when(mPreferenceScreen).findPreference(anyString());
+
+        mButtonOn = new Button(mContext);
+        mButtonOn.setId(R.id.state_on_button);
+        doReturn(mButtonOn).when(mPreference).getStateOnButton();
+        mButtonOff = new Button(mContext);
+        mButtonOff.setId(R.id.state_off_button);
+        doReturn(mButtonOff).when(mPreference).getStateOffButton();
+
+        mController = new BatterySaverButtonPreferenceController(mContext, mLifecycle);
+        mController.displayPreference(mPreferenceScreen);
+    }
+
+    @Test
+    public void testUpdateState_lowPowerOn_displayButtonOff() {
+        mShadowPowerManager.setIsPowerSaveMode(true);
+
+        mController.updateState(mPreference);
+
+        assertThat(mButtonOn.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mButtonOff.getVisibility()).isEqualTo(View.VISIBLE);
+    }
+
+    @Test
+    public void testUpdateState_lowPowerOff_displayButtonOn() {
+        mShadowPowerManager.setIsPowerSaveMode(false);
+
+        mController.updateState(mPreference);
+
+        assertThat(mButtonOn.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mButtonOff.getVisibility()).isEqualTo(View.GONE);
+    }
+
+    @Test
+    public void testOnClick_clickButtonOn_setPowerSaveMode() {
+        mController.onClick(mButtonOn);
+
+        assertThat(mShadowPowerManager.isPowerSaveMode()).isTrue();
+    }
+
+    @Test
+    public void testOnClick_clickButtonOff_clearPowerSaveMode() {
+        mController.onClick(mButtonOff);
+
+        assertThat(mShadowPowerManager.isPowerSaveMode()).isFalse();
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowPowerManager.java b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowPowerManager.java
new file mode 100644
index 00000000000..3bc807adc17
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowPowerManager.java
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.settings.testutils.shadow;
+
+import android.os.PowerManager;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+
+@Implements(PowerManager.class)
+public class ShadowPowerManager extends org.robolectric.shadows.ShadowPowerManager {
+    @Implementation
+    public boolean setPowerSaveMode(boolean mode) {
+        setIsPowerSaveMode(mode);
+        return true;
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/widget/TwoStateButtonPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/widget/TwoStateButtonPreferenceControllerTest.java
new file mode 100644
index 00000000000..b1608889935
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/widget/TwoStateButtonPreferenceControllerTest.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.settings.widget;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+
+import android.content.Context;
+import android.support.v7.preference.PreferenceScreen;
+import android.view.View;
+import android.widget.Button;
+
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
+
+@RunWith(SettingsRobolectricTestRunner.class)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+public class TwoStateButtonPreferenceControllerTest {
+    private static final String KEY = "pref_key";
+
+    @Mock
+    private PreferenceScreen mPreferenceScreen;
+    @Mock
+    private TwoStateButtonPreference mPreference;
+    private TwoStateButtonPreferenceController mController;
+    private Context mContext;
+    private Button mButtonOn;
+    private Button mButtonOff;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = spy(RuntimeEnvironment.application);
+        doReturn(mPreference).when(mPreferenceScreen).findPreference(anyString());
+        mButtonOn = new Button(mContext);
+        doReturn(mButtonOn).when(mPreference).getStateOnButton();
+        mButtonOff = new Button(mContext);
+        doReturn(mButtonOff).when(mPreference).getStateOffButton();
+
+        mController = new TestButtonsPreferenceController(mContext, KEY);
+        mController.displayPreference(mPreferenceScreen);
+    }
+
+    @Test
+    public void testUpdateButtons_stateOn_onlyShowButtonOn() {
+        mController.updateButton(true /* stateOn */);
+
+        assertThat(mButtonOn.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mButtonOff.getVisibility()).isEqualTo(View.GONE);
+    }
+
+    @Test
+    public void testUpdateButtons_stateOff_onlyShowButtonOff() {
+        mController.updateButton(false /* stateOn */);
+
+        assertThat(mButtonOn.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mButtonOff.getVisibility()).isEqualTo(View.VISIBLE);
+    }
+
+    /**
+     * Controller to test methods in {@link TwoStateButtonPreferenceController}
+     */
+    public static class TestButtonsPreferenceController extends
+            TwoStateButtonPreferenceController {
+
+        TestButtonsPreferenceController(Context context, String key) {
+            super(context, key);
+        }
+
+        @Override
+        public void onButtonClicked(boolean stateOn) {
+            //do nothing
+        }
+
+        @Override
+        public int getAvailabilityStatus() {
+            return AVAILABLE;
+        }
+    }
+}