From 10b7a7d48369d22d21fe4eff7d037b4a7c52ec4a Mon Sep 17 00:00:00 2001
From: Stanley Wang <stanleytfwang@google.com>
Date: Fri, 29 Jan 2021 01:17:43 +0800
Subject: [PATCH] Use MainSwitchPreference to replace the button style
 switches.

Fix: 177967926
Fix: 177968295
Fix: 177967925
Fix: 177968078

Test: Run robotest and apply the widget in Settings and see the ui
Change-Id: Ie854de96e5495fa564fb8a097ed4547bbd2b10c5
---
 res/layout/dark_ui_activation_button.xml      |  40 -------
 .../night_display_activation_button.xml       |  40 -------
 res/layout/zen_mode_settings_button.xml       |  55 ----------
 res/values/strings.xml                        |  12 +++
 res/xml/auto_brightness_detail.xml            |  17 ++-
 res/xml/dark_mode_settings.xml                |  13 +--
 res/xml/nfc_and_payment_settings.xml          |  11 +-
 res/xml/night_display_settings.xml            |  14 +--
 res/xml/smart_battery_detail.xml              |  14 ++-
 res/xml/zen_mode_settings.xml                 |   7 +-
 ...oBrightnessDetailPreferenceController.java |  29 ++++-
 ...DisplayActivationPreferenceController.java |  93 +++++-----------
 ...arkModeActivationPreferenceController.java |  80 +++-----------
 .../SmartBatteryPreferenceController.java     |  24 +++--
 src/com/android/settings/nfc/NfcEnabler.java  |  15 +--
 .../settings/nfc/NfcPreferenceController.java |  33 +++---
 .../app/BlockPreferenceController.java        |   6 +-
 .../ZenModeButtonPreferenceController.java    | 101 ++++++++----------
 .../widget/SettingsMainSwitchPreference.java  |  28 +++--
 ...layActivationPreferenceControllerTest.java |  20 ++--
 ...odeActivationPreferenceControllerTest.java |  71 +++++-------
 .../SmartBatteryPreferenceControllerTest.java |  22 ++--
 .../android/settings/nfc/NfcEnablerTest.java  |   4 +-
 .../nfc/NfcPreferenceControllerTest.java      |   6 +-
 ...ZenModeButtonPreferenceControllerTest.java |  65 +++++------
 25 files changed, 307 insertions(+), 513 deletions(-)
 delete mode 100644 res/layout/dark_ui_activation_button.xml
 delete mode 100644 res/layout/night_display_activation_button.xml
 delete mode 100644 res/layout/zen_mode_settings_button.xml

diff --git a/res/layout/dark_ui_activation_button.xml b/res/layout/dark_ui_activation_button.xml
deleted file mode 100644
index 5f9eefcded2..00000000000
--- a/res/layout/dark_ui_activation_button.xml
+++ /dev/null
@@ -1,40 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  Copyright (C) 2019 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.
-  -->
-
-<LinearLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:minHeight="?android:attr/listPreferredItemHeight">
-
-    <Button
-        android:id="@+id/dark_ui_turn_on_button"
-        style="@style/ActionPrimaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_vertical" />
-
-    <Button
-        android:id="@+id/dark_ui_turn_off_button"
-        style="@style/ActionSecondaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_vertical" />
-
-</LinearLayout>
\ No newline at end of file
diff --git a/res/layout/night_display_activation_button.xml b/res/layout/night_display_activation_button.xml
deleted file mode 100644
index df160ec33b8..00000000000
--- a/res/layout/night_display_activation_button.xml
+++ /dev/null
@@ -1,40 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  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.
-  -->
-
-<LinearLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:minHeight="?android:attr/listPreferredItemHeight">
-
-    <Button
-        android:id="@+id/night_display_turn_on_button"
-        style="@style/ActionPrimaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_vertical" />
-
-    <Button
-        android:id="@+id/night_display_turn_off_button"
-        style="@style/ActionSecondaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_vertical" />
-
-</LinearLayout>
\ No newline at end of file
diff --git a/res/layout/zen_mode_settings_button.xml b/res/layout/zen_mode_settings_button.xml
deleted file mode 100644
index db889ff9300..00000000000
--- a/res/layout/zen_mode_settings_button.xml
+++ /dev/null
@@ -1,55 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  Copyright (C) 2017 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.
-  -->
-
-<LinearLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:orientation="vertical">
-
-    <TextView
-        android:text="@string/zen_mode_settings_summary"
-        android:textAppearance="?android:attr/textAppearanceSmall"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_marginEnd="@dimen/description_margin_sides"
-        android:layout_marginTop="@dimen/zen_mode_settings_button_margin_vertical"/>
-
-    <Button
-        android:id="@+id/zen_mode_settings_turn_on_button"
-        style="@style/ActionPrimaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_marginEnd="@dimen/description_margin_sides"
-        android:layout_marginTop="@dimen/zen_mode_settings_button_margin_vertical"
-        android:layout_marginBottom="@dimen/zen_mode_settings_button_margin_vertical"
-        android:text="@string/zen_mode_button_turn_on"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"/>
-
-    <Button
-        android:id="@+id/zen_mode_settings_turn_off_button"
-        style="@style/ActionSecondaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:layout_marginEnd="@dimen/description_margin_sides"
-        android:layout_marginTop="@dimen/zen_mode_settings_button_margin_vertical"
-        android:layout_marginBottom="@dimen/zen_mode_settings_button_margin_vertical"
-        android:text="@string/zen_mode_button_turn_off"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"/>
-
-</LinearLayout>
\ No newline at end of file
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 09aea0e1f0b..8a811c0c08a 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -12670,6 +12670,18 @@
     <string name="graphics_driver_main_switch_title" translatable="false">Use graphics driver preferences</string>
     <!-- Do not translate. Title for battery saver main switch preferences. [CHAR LIMIT=50] -->
     <string name="battery_saver_main_switch_title" translatable="false">Use battery saver</string>
+    <!-- Do not translate. Title for Do Not Disturb main switch preferences. [CHAR LIMIT=50] -->
+    <string name="do_not_disturb_main_switch_title" translatable="false">Use Do Not Disturb</string>
+    <!-- Do not translate. Title for Dark theme main switch preferences. [CHAR LIMIT=50] -->
+    <string name="dark_theme_main_switch_title" translatable="false">Use Dark Theme</string>
+    <!-- Do not translate. Title for Night Light main switch preferences. [CHAR LIMIT=50] -->
+    <string name="night_light_main_switch_title" translatable="false">Use Night Light</string>
+    <!-- Do not translate. Title for NFC main switch preferences. [CHAR LIMIT=50] -->
+    <string name="nfc_main_switch_title" translatable="false">Use NFC</string>
+    <!-- Do not translate. Title for Adaptive Battery main switch preferences. [CHAR LIMIT=50] -->
+    <string name="adaptive_battery_main_switch_title" translatable="false">Use adaptive battery</string>
+    <!-- Do not translate. Title for adaptive brightness main switch preferences. [CHAR LIMIT=50] -->
+    <string name="adaptive_brightness_main_switch_title" translatable="false">Use adaptive brightness</string>
 
     <!-- Default preference title for showing all apps on device [CHAR_LIMIT=50]-->
     <string name="default_see_all_apps_title">See all apps</string>
diff --git a/res/xml/auto_brightness_detail.xml b/res/xml/auto_brightness_detail.xml
index 6b304efcaf7..ebcf1741b77 100644
--- a/res/xml/auto_brightness_detail.xml
+++ b/res/xml/auto_brightness_detail.xml
@@ -20,6 +20,12 @@
     xmlns:settings="http://schemas.android.com/apk/res-auto"
     android:title="@string/auto_brightness_title">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="auto_brightness"
+        android:title="@string/adaptive_brightness_main_switch_title"
+        settings:keywords="@string/keywords_display_auto_brightness"
+        settings:controller="com.android.settings.display.AutoBrightnessDetailPreferenceController"/>
+
     <com.android.settings.widget.VideoPreference
         android:key="auto_brightness_video"
         android:title="@string/summary_placeholder"
@@ -27,17 +33,6 @@
         settings:preview="@drawable/aab_brightness"
         settings:controller="com.android.settings.widget.VideoPreferenceController"/>
 
-    <!-- Cross-listed item, if you change this, also change it in power_usage_summary.xml -->
-    <com.android.settingslib.RestrictedSwitchPreference
-        android:key="auto_brightness"
-        android:title="@string/auto_brightness_title"
-        android:summary="@string/auto_brightness_summary"
-        settings:keywords="@string/keywords_display_auto_brightness"
-        settings:controller="com.android.settings.display.AutoBrightnessDetailPreferenceController"
-        settings:useAdminDisabledSummary="true"
-        settings:userRestriction="no_config_brightness"
-        settings:allowDividerAbove="true" />
-
     <com.android.settingslib.widget.FooterPreference
         android:key="auto_brightness_footer"
         android:title="@string/auto_brightness_description"
diff --git a/res/xml/dark_mode_settings.xml b/res/xml/dark_mode_settings.xml
index 19279e0c6d0..5ac9e261b62 100644
--- a/res/xml/dark_mode_settings.xml
+++ b/res/xml/dark_mode_settings.xml
@@ -20,6 +20,11 @@
     android:title="@string/dark_ui_mode"
     settings:keywords="@string/keywords_dark_ui_mode">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="dark_ui_activated"
+        android:title="@string/dark_theme_main_switch_title"
+        settings:controller="com.android.settings.display.darkmode.DarkModeActivationPreferenceController"/>
+
     <com.android.settingslib.widget.BannerMessagePreference
         android:key="dark_ui_location_off"
         android:title="@string/twilight_mode_location_off_dialog_message"
@@ -50,12 +55,4 @@
         android:title="@string/night_display_end_time_title"
         settings:searchable="false"/>
 
-    <com.android.settingslib.widget.LayoutPreference
-        android:key="dark_ui_activated"
-        android:title="@string/dark_ui_title"
-        android:selectable="false"
-        android:layout="@layout/dark_ui_activation_button"
-        settings:allowDividerBelow="true"
-        settings:controller="com.android.settings.display.darkmode.DarkModeActivationPreferenceController"/>
-
 </PreferenceScreen>
diff --git a/res/xml/nfc_and_payment_settings.xml b/res/xml/nfc_and_payment_settings.xml
index 8ff4983579b..09ce27769ab 100644
--- a/res/xml/nfc_and_payment_settings.xml
+++ b/res/xml/nfc_and_payment_settings.xml
@@ -19,18 +19,17 @@
     xmlns:settings="http://schemas.android.com/apk/res-auto"
     android:title="@string/nfc_quick_toggle_title">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="toggle_nfc"
+        android:title="@string/nfc_main_switch_title"
+        settings:controller="com.android.settings.nfc.NfcPreferenceController"/>
+
     <com.android.settingslib.widget.LayoutPreference
         android:key="nfc_detection_point"
         android:selectable="false"
         android:layout="@layout/nfc_detection_point"
         settings:controller="com.android.settings.nfc.NfcDetectionPointController"/>
 
-    <SwitchPreference
-        android:key="toggle_nfc"
-        android:title="@string/nfc_quick_toggle_title"
-        android:summary="@string/nfc_quick_toggle_summary"
-        settings:controller="com.android.settings.nfc.NfcPreferenceController"/>
-
     <SwitchPreference
         android:key="nfc_secure_settings"
         android:title="@string/nfc_secure_settings_title"
diff --git a/res/xml/night_display_settings.xml b/res/xml/night_display_settings.xml
index 97f5dd27627..23afccf48fe 100644
--- a/res/xml/night_display_settings.xml
+++ b/res/xml/night_display_settings.xml
@@ -20,6 +20,11 @@
     android:title="@string/night_display_title"
     settings:keywords="@string/keywords_display_night_display">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="night_display_activated"
+        android:title="@string/night_light_main_switch_title"
+        settings:controller="com.android.settings.display.NightDisplayActivationPreferenceController"/>
+
     <com.android.settingslib.widget.BannerMessagePreference
         android:key="dark_ui_location_off"
         android:title="@string/twilight_mode_location_off_dialog_message"
@@ -54,13 +59,4 @@
         settings:controller="com.android.settings.display.NightDisplayIntensityPreferenceController"
         settings:unavailableSliceSubtitle="@string/night_display_not_currently_on"/>
 
-    <com.android.settingslib.widget.LayoutPreference
-        android:key="night_display_activated"
-        android:title="@string/night_display_title"
-        android:selectable="false"
-        android:layout="@layout/night_display_activation_button"
-        settings:allowDividerBelow="true"
-        settings:keywords="@string/keywords_display_night_display"
-        settings:controller="com.android.settings.display.NightDisplayActivationPreferenceController"/>
-
 </PreferenceScreen>
diff --git a/res/xml/smart_battery_detail.xml b/res/xml/smart_battery_detail.xml
index 87eeb8621ba..2d361191810 100644
--- a/res/xml/smart_battery_detail.xml
+++ b/res/xml/smart_battery_detail.xml
@@ -22,6 +22,11 @@
     android:title="@string/smart_battery_manager_title"
     settings:searchable="false">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="smart_battery"
+        android:title="@string/adaptive_battery_main_switch_title"
+        settings:controller="com.android.settings.fuelgauge.SmartBatteryPreferenceController"/>
+
     <com.android.settings.widget.VideoPreference
         android:key="auto_awesome_battery"
         android:title="@string/summary_placeholder"
@@ -29,13 +34,6 @@
         settings:preview="@drawable/auto_awesome_battery"
         settings:controller="com.android.settings.widget.VideoPreferenceController"/>
 
-    <SwitchPreference
-        android:key="smart_battery"
-        android:title="@string/smart_battery_title"
-        android:summary="@string/smart_battery_summary"
-        settings:controller="com.android.settings.fuelgauge.SmartBatteryPreferenceController"
-        settings:allowDividerAbove="true"/>
-
     <SwitchPreference
         android:key="auto_restriction"
         android:title="@string/battery_auto_restriction_title"
@@ -51,4 +49,4 @@
         android:title="@string/smart_battery_footer"
         android:selectable="false"
         settings:searchable="false"/>
-</PreferenceScreen>
\ No newline at end of file
+</PreferenceScreen>
diff --git a/res/xml/zen_mode_settings.xml b/res/xml/zen_mode_settings.xml
index 6764e288d43..78dee028666 100644
--- a/res/xml/zen_mode_settings.xml
+++ b/res/xml/zen_mode_settings.xml
@@ -21,12 +21,9 @@
     android:title="@string/zen_mode_settings_title">
 
     <!-- Turn on DND button -->
-    <com.android.settingslib.widget.LayoutPreference
+    <com.android.settingslib.widget.MainSwitchPreference
         android:key="zen_mode_toggle"
-        android:title="@string/zen_mode_settings_title"
-        android:selectable="false"
-        android:layout="@layout/zen_mode_settings_button"
-        settings:allowDividerBelow="true"
+        android:title="@string/do_not_disturb_main_switch_title"
         settings:keywords="@string/keywords_zen_mode_settings"/>
 
     <PreferenceCategory
diff --git a/src/com/android/settings/display/AutoBrightnessDetailPreferenceController.java b/src/com/android/settings/display/AutoBrightnessDetailPreferenceController.java
index 030c5b92f09..3138ae09b86 100644
--- a/src/com/android/settings/display/AutoBrightnessDetailPreferenceController.java
+++ b/src/com/android/settings/display/AutoBrightnessDetailPreferenceController.java
@@ -17,8 +17,18 @@
 package com.android.settings.display;
 
 import android.content.Context;
+import android.widget.Switch;
 
-public class AutoBrightnessDetailPreferenceController extends AutoBrightnessPreferenceController {
+import androidx.preference.PreferenceScreen;
+
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
+
+/**
+ * Controller that updates the adaptive brightness.
+ */
+public class AutoBrightnessDetailPreferenceController extends
+        AutoBrightnessPreferenceController implements OnMainSwitchChangeListener {
 
     public AutoBrightnessDetailPreferenceController(Context context, String key) {
         super(context, key);
@@ -36,4 +46,21 @@ public class AutoBrightnessDetailPreferenceController extends AutoBrightnessPref
     public boolean isPublicSlice() {
         return true;
     }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+
+        MainSwitchPreference pref = (MainSwitchPreference) screen.findPreference(
+                getPreferenceKey());
+        pref.addOnSwitchChangeListener(this);
+        pref.updateStatus(isChecked());
+    }
+
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        if (isChecked != isChecked()) {
+            setChecked(isChecked);
+        }
+    }
 }
diff --git a/src/com/android/settings/display/NightDisplayActivationPreferenceController.java b/src/com/android/settings/display/NightDisplayActivationPreferenceController.java
index cf9332bfcc3..e6217ff7eb2 100644
--- a/src/com/android/settings/display/NightDisplayActivationPreferenceController.java
+++ b/src/com/android/settings/display/NightDisplayActivationPreferenceController.java
@@ -16,43 +16,30 @@
 
 package com.android.settings.display;
 
-import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_FOCUSED;
-
 import android.content.Context;
 import android.hardware.display.ColorDisplayManager;
 import android.text.TextUtils;
-import android.view.View;
-import android.view.View.OnClickListener;
-import android.widget.Button;
+import android.widget.Switch;
 
 import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
 
-import com.android.settings.R;
 import com.android.settings.core.TogglePreferenceController;
 import com.android.settings.overlay.FeatureFactory;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
-import com.android.settingslib.widget.LayoutPreference;
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
 
-public class NightDisplayActivationPreferenceController extends TogglePreferenceController {
+/**
+ * Controller that updates the night display.
+ */
+public class NightDisplayActivationPreferenceController extends
+        TogglePreferenceController implements OnMainSwitchChangeListener {
 
     private final MetricsFeatureProvider mMetricsFeatureProvider;
     private ColorDisplayManager mColorDisplayManager;
     private NightDisplayTimeFormatter mTimeFormatter;
-    private LayoutPreference mPreference;
-
-    private Button mTurnOffButton;
-    private Button mTurnOnButton;
-
-    private final OnClickListener mListener = new OnClickListener() {
-        @Override
-        public void onClick(View v) {
-            mMetricsFeatureProvider.logClickedPreference(mPreference, getMetricsCategory());
-            mColorDisplayManager.setNightDisplayActivated(
-                    !mColorDisplayManager.isNightDisplayActivated());
-            updateStateInternal(true);
-        }
-    };
+    private MainSwitchPreference mPreference;
 
     public NightDisplayActivationPreferenceController(Context context, String key) {
         super(context, key);
@@ -64,7 +51,7 @@ public class NightDisplayActivationPreferenceController extends TogglePreference
 
     @Override
     public int getAvailabilityStatus() {
-        return ColorDisplayManager.isNightDisplayAvailable(mContext) ? AVAILABLE
+        return ColorDisplayManager.isNightDisplayAvailable(mContext) ? AVAILABLE_UNSEARCHABLE
                 : UNSUPPORTED_ON_DEVICE;
     }
 
@@ -82,16 +69,24 @@ public class NightDisplayActivationPreferenceController extends TogglePreference
     public void displayPreference(PreferenceScreen screen) {
         super.displayPreference(screen);
 
-        mPreference = screen.findPreference(getPreferenceKey());
-        mTurnOnButton = mPreference.findViewById(R.id.night_display_turn_on_button);
-        mTurnOnButton.setOnClickListener(mListener);
-        mTurnOffButton = mPreference.findViewById(R.id.night_display_turn_off_button);
-        mTurnOffButton.setOnClickListener(mListener);
+        mPreference = (MainSwitchPreference) screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
+        mPreference.updateStatus(mColorDisplayManager.isNightDisplayActivated());
+    }
+
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        final boolean activated = mColorDisplayManager.isNightDisplayActivated();
+        if (isChecked != activated) {
+            // TODO(b/179017365): Create a controller which extends TogglePreferenceController to
+            //  control the toggle preference.
+            setChecked(isChecked);
+        }
     }
 
     @Override
     public final void updateState(Preference preference) {
-        updateStateInternal(false);
+        updateStateInternal();
     }
 
     /** FOR SLICES */
@@ -111,46 +106,14 @@ public class NightDisplayActivationPreferenceController extends TogglePreference
         return mTimeFormatter.getAutoModeSummary(mContext, mColorDisplayManager);
     }
 
-    private void updateStateInternal(boolean selfChanged) {
-        if (mTurnOnButton == null || mTurnOffButton == null) {
-            return;
-        }
-
+    private void updateStateInternal() {
         final boolean isActivated = mColorDisplayManager.isNightDisplayActivated();
         final int autoMode = mColorDisplayManager.getNightDisplayAutoMode();
 
-        String buttonText;
         if (autoMode == ColorDisplayManager.AUTO_MODE_CUSTOM_TIME) {
-            buttonText = mContext.getString(isActivated
-                            ? R.string.night_display_activation_off_custom
-                            : R.string.night_display_activation_on_custom,
-                    mTimeFormatter.getFormattedTimeString(isActivated
-                            ? mColorDisplayManager.getNightDisplayCustomStartTime()
-                            : mColorDisplayManager.getNightDisplayCustomEndTime()));
-        } else if (autoMode == ColorDisplayManager.AUTO_MODE_TWILIGHT) {
-            buttonText = mContext.getString(isActivated
-                    ? R.string.night_display_activation_off_twilight
-                    : R.string.night_display_activation_on_twilight);
-        } else {
-            buttonText = mContext.getString(isActivated
-                    ? R.string.night_display_activation_off_manual
-                    : R.string.night_display_activation_on_manual);
-        }
-
-        if (isActivated) {
-            mTurnOnButton.setVisibility(View.GONE);
-            mTurnOffButton.setVisibility(View.VISIBLE);
-            mTurnOffButton.setText(buttonText);
-            if (selfChanged) {
-                mTurnOffButton.sendAccessibilityEvent(TYPE_VIEW_FOCUSED);
-            }
-        } else {
-            mTurnOnButton.setVisibility(View.VISIBLE);
-            mTurnOffButton.setVisibility(View.GONE);
-            mTurnOnButton.setText(buttonText);
-            if (selfChanged) {
-                mTurnOnButton.sendAccessibilityEvent(TYPE_VIEW_FOCUSED);
-            }
+            mTimeFormatter.getFormattedTimeString(isActivated
+                    ? mColorDisplayManager.getNightDisplayCustomStartTime()
+                    : mColorDisplayManager.getNightDisplayCustomEndTime());
         }
     }
 }
diff --git a/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceController.java b/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceController.java
index 1027dac4e87..800e7e506b0 100644
--- a/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceController.java
+++ b/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceController.java
@@ -19,8 +19,7 @@ import android.app.UiModeManager;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.os.PowerManager;
-import android.view.View;
-import android.widget.Button;
+import android.widget.Switch;
 
 import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
@@ -29,22 +28,22 @@ import com.android.settings.R;
 import com.android.settings.core.BasePreferenceController;
 import com.android.settings.overlay.FeatureFactory;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
-import com.android.settingslib.widget.LayoutPreference;
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
 
 import java.time.LocalTime;
 
 /**
  * Controller for activate/deactivate night mode button
  */
-public class DarkModeActivationPreferenceController extends BasePreferenceController {
+public class DarkModeActivationPreferenceController extends BasePreferenceController implements
+        OnMainSwitchChangeListener {
 
     private final UiModeManager mUiModeManager;
     private final MetricsFeatureProvider mMetricsFeatureProvider;
     private PowerManager mPowerManager;
-    private Button mTurnOffButton;
-    private Button mTurnOnButton;
     private TimeFormatter mFormat;
-    private LayoutPreference mPreference;
+    private MainSwitchPreference mPreference;
 
     public DarkModeActivationPreferenceController(Context context, String preferenceKey) {
         super(context, preferenceKey);
@@ -62,49 +61,9 @@ public class DarkModeActivationPreferenceController extends BasePreferenceContro
 
     @Override
     public final void updateState(Preference preference) {
-
-        final boolean batterySaver = mPowerManager.isPowerSaveMode();
-        if (batterySaver) {
-            mTurnOnButton.setVisibility(View.GONE);
-            mTurnOffButton.setVisibility(View.GONE);
-            return;
-        }
-
         final boolean active = (mContext.getResources().getConfiguration().uiMode
                 & Configuration.UI_MODE_NIGHT_YES) != 0;
-        updateNightMode(active);
-    }
-
-    private void updateNightMode(boolean active) {
-        final int mode = mUiModeManager.getNightMode();
-        String buttonText;
-
-        if (mode == UiModeManager.MODE_NIGHT_AUTO) {
-            buttonText = mContext.getString(active
-                    ? R.string.dark_ui_activation_off_auto
-                    : R.string.dark_ui_activation_on_auto);
-        } else if (mode == UiModeManager.MODE_NIGHT_CUSTOM) {
-            final LocalTime time = active
-                    ? mUiModeManager.getCustomNightModeStart()
-                    : mUiModeManager.getCustomNightModeEnd();
-            final String timeStr = mFormat.of(time);
-            buttonText = mContext.getString(active
-                    ? R.string.dark_ui_activation_off_custom
-                    : R.string.dark_ui_activation_on_custom, timeStr);
-        } else {
-            buttonText = mContext.getString(active
-                    ? R.string.dark_ui_activation_off_manual
-                    : R.string.dark_ui_activation_on_manual);
-        }
-        if (active) {
-            mTurnOnButton.setVisibility(View.GONE);
-            mTurnOffButton.setVisibility(View.VISIBLE);
-            mTurnOffButton.setText(buttonText);
-        } else {
-            mTurnOnButton.setVisibility(View.VISIBLE);
-            mTurnOffButton.setVisibility(View.GONE);
-            mTurnOnButton.setText(buttonText);
-        }
+        mPreference.updateStatus(active);
     }
 
     @Override
@@ -132,26 +91,19 @@ public class DarkModeActivationPreferenceController extends BasePreferenceContro
         }
     }
 
-    private final View.OnClickListener mListener = new View.OnClickListener() {
-        @Override
-        public void onClick(View v) {
-            mMetricsFeatureProvider.logClickedPreference(mPreference, getMetricsCategory());
-            final boolean active = (mContext.getResources().getConfiguration().uiMode
-                    & Configuration.UI_MODE_NIGHT_YES) != 0;
-            mUiModeManager.setNightModeActivated(!active);
-            updateNightMode(!active);
-        }
-    };
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        mMetricsFeatureProvider.logClickedPreference(mPreference, getMetricsCategory());
+        final boolean active = (mContext.getResources().getConfiguration().uiMode
+                & Configuration.UI_MODE_NIGHT_YES) != 0;
+        mUiModeManager.setNightModeActivated(!active);
+    }
 
     @Override
     public void displayPreference(PreferenceScreen screen) {
         super.displayPreference(screen);
-
-        mPreference = screen.findPreference(getPreferenceKey());
-        mTurnOnButton = mPreference.findViewById(R.id.dark_ui_turn_on_button);
-        mTurnOnButton.setOnClickListener(mListener);
-        mTurnOffButton = mPreference.findViewById(R.id.dark_ui_turn_off_button);
-        mTurnOffButton.setOnClickListener(mListener);
+        mPreference = (MainSwitchPreference) screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
     }
 
     @Override
diff --git a/src/com/android/settings/fuelgauge/SmartBatteryPreferenceController.java b/src/com/android/settings/fuelgauge/SmartBatteryPreferenceController.java
index 040af9f7d20..ea277a12add 100644
--- a/src/com/android/settings/fuelgauge/SmartBatteryPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/SmartBatteryPreferenceController.java
@@ -20,22 +20,27 @@ package com.android.settings.fuelgauge;
 import android.content.Context;
 import android.provider.Settings;
 import android.text.TextUtils;
+import android.widget.Switch;
 
 import androidx.preference.Preference;
-import androidx.preference.SwitchPreference;
+import androidx.preference.PreferenceScreen;
 
 import com.android.settings.core.BasePreferenceController;
 import com.android.settings.overlay.FeatureFactory;
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
 
 /**
  * Controller to change and update the smart battery toggle
  */
 public class SmartBatteryPreferenceController extends BasePreferenceController implements
-        Preference.OnPreferenceChangeListener {
+        OnMainSwitchChangeListener {
+
     private static final String KEY_SMART_BATTERY = "smart_battery";
     private static final int ON = 1;
     private static final int OFF = 0;
     private PowerUsageFeatureProvider mPowerUsageFeatureProvider;
+    private MainSwitchPreference mPreference;
 
     public SmartBatteryPreferenceController(Context context) {
         super(context, KEY_SMART_BATTERY);
@@ -65,14 +70,19 @@ public class SmartBatteryPreferenceController extends BasePreferenceController i
         super.updateState(preference);
         final boolean smartBatteryOn = Settings.Global.getInt(mContext.getContentResolver(),
                 Settings.Global.ADAPTIVE_BATTERY_MANAGEMENT_ENABLED, ON) == ON;
-        ((SwitchPreference) preference).setChecked(smartBatteryOn);
+        ((MainSwitchPreference) preference).updateStatus(smartBatteryOn);
     }
 
     @Override
-    public boolean onPreferenceChange(Preference preference, Object newValue) {
-        final boolean smartBatteryOn = (Boolean) newValue;
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        mPreference = (MainSwitchPreference) screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
+    }
+
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
         Settings.Global.putInt(mContext.getContentResolver(),
-                Settings.Global.ADAPTIVE_BATTERY_MANAGEMENT_ENABLED, smartBatteryOn ? ON : OFF);
-        return true;
+                Settings.Global.ADAPTIVE_BATTERY_MANAGEMENT_ENABLED, isChecked ? ON : OFF);
     }
 }
diff --git a/src/com/android/settings/nfc/NfcEnabler.java b/src/com/android/settings/nfc/NfcEnabler.java
index 777e7d114bd..88482d59d8b 100644
--- a/src/com/android/settings/nfc/NfcEnabler.java
+++ b/src/com/android/settings/nfc/NfcEnabler.java
@@ -21,16 +21,17 @@ import android.nfc.NfcAdapter;
 import android.provider.Settings;
 
 import androidx.annotation.VisibleForTesting;
-import androidx.preference.SwitchPreference;
+
+import com.android.settingslib.widget.MainSwitchPreference;
 
 /**
  * NfcEnabler is a helper to manage the Nfc on/off checkbox preference. It turns on/off Nfc
  * and ensures the summary of the preference reflects the current state.
  */
 public class NfcEnabler extends BaseNfcEnabler {
-    private final SwitchPreference mPreference;
+    private final MainSwitchPreference mPreference;
 
-    public NfcEnabler(Context context, SwitchPreference preference) {
+    public NfcEnabler(Context context, MainSwitchPreference preference) {
         super(context);
         mPreference = preference;
     }
@@ -39,19 +40,19 @@ public class NfcEnabler extends BaseNfcEnabler {
     protected void handleNfcStateChanged(int newState) {
         switch (newState) {
             case NfcAdapter.STATE_OFF:
-                mPreference.setChecked(false);
+                mPreference.updateStatus(false);
                 mPreference.setEnabled(isToggleable());
                 break;
             case NfcAdapter.STATE_ON:
-                mPreference.setChecked(true);
+                mPreference.updateStatus(true);
                 mPreference.setEnabled(true);
                 break;
             case NfcAdapter.STATE_TURNING_ON:
-                mPreference.setChecked(true);
+                mPreference.updateStatus(true);
                 mPreference.setEnabled(false);
                 break;
             case NfcAdapter.STATE_TURNING_OFF:
-                mPreference.setChecked(false);
+                mPreference.updateStatus(false);
                 mPreference.setEnabled(false);
                 break;
         }
diff --git a/src/com/android/settings/nfc/NfcPreferenceController.java b/src/com/android/settings/nfc/NfcPreferenceController.java
index 9bf013c8c3f..b30e7c3fe5d 100644
--- a/src/com/android/settings/nfc/NfcPreferenceController.java
+++ b/src/com/android/settings/nfc/NfcPreferenceController.java
@@ -23,24 +23,27 @@ import android.net.Uri;
 import android.nfc.NfcAdapter;
 import android.provider.Settings;
 import android.util.Log;
+import android.widget.Switch;
 
 import androidx.preference.PreferenceScreen;
-import androidx.preference.SwitchPreference;
 
 import com.android.settings.core.TogglePreferenceController;
 import com.android.settings.slices.SliceBackgroundWorker;
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
 import com.android.settingslib.core.lifecycle.events.OnPause;
 import com.android.settingslib.core.lifecycle.events.OnResume;
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
 
 import java.io.IOException;
 
 public class NfcPreferenceController extends TogglePreferenceController
-        implements LifecycleObserver, OnResume, OnPause {
+        implements LifecycleObserver, OnResume, OnPause, OnMainSwitchChangeListener {
 
     public static final String KEY_TOGGLE_NFC = "toggle_nfc";
     private final NfcAdapter mNfcAdapter;
     private NfcEnabler mNfcEnabler;
+    private MainSwitchPreference mPreference;
 
     public NfcPreferenceController(Context context, String key) {
         super(context, key);
@@ -55,10 +58,16 @@ public class NfcPreferenceController extends TogglePreferenceController
             return;
         }
 
-        final SwitchPreference switchPreference = screen.findPreference(getPreferenceKey());
-
-        mNfcEnabler = new NfcEnabler(mContext, switchPreference);
+        mPreference = (MainSwitchPreference) screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
+        mNfcEnabler = new NfcEnabler(mContext, mPreference);
+    }
 
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        if (isChecked != mNfcAdapter.isEnabled()) {
+            setChecked(isChecked);
+        }
     }
 
     @Override
@@ -129,12 +138,12 @@ public class NfcPreferenceController extends TogglePreferenceController
      * Listener for background changes to NFC.
      *
      * <p>
-     *     Listen to broadcasts from {@link NfcAdapter}. The worker will call notify changed on the
-     *     NFC Slice only when the following extras are present in the broadcast:
-     *     <ul>
-     *      <li>{@link NfcAdapter#STATE_ON}</li>
-     *      <li>{@link NfcAdapter#STATE_OFF}</li>
-     *     </ul>
+     * Listen to broadcasts from {@link NfcAdapter}. The worker will call notify changed on the
+     * NFC Slice only when the following extras are present in the broadcast:
+     * <ul>
+     * <li>{@link NfcAdapter#STATE_ON}</li>
+     * <li>{@link NfcAdapter#STATE_OFF}</li>
+     * </ul>
      */
     public static class NfcSliceWorker extends SliceBackgroundWorker<Void> {
 
@@ -185,7 +194,7 @@ public class NfcPreferenceController extends TogglePreferenceController
                         NO_EXTRA);
 
                 // Do nothing if state change is empty, or an intermediate step.
-                if ( (nfcStateExtra == NO_EXTRA)
+                if ((nfcStateExtra == NO_EXTRA)
                         || (nfcStateExtra == NfcAdapter.STATE_TURNING_ON)
                         || (nfcStateExtra == NfcAdapter.STATE_TURNING_OFF)) {
                     Log.d(TAG, "Transitional update, dropping broadcast");
diff --git a/src/com/android/settings/notification/app/BlockPreferenceController.java b/src/com/android/settings/notification/app/BlockPreferenceController.java
index e5a7e980c5f..4df3e281f61 100644
--- a/src/com/android/settings/notification/app/BlockPreferenceController.java
+++ b/src/com/android/settings/notification/app/BlockPreferenceController.java
@@ -72,16 +72,16 @@ public class BlockPreferenceController extends NotificationPreferenceController
             bar.setDisabledByAdmin(mAdmin);
 
             if (mChannel != null && !isChannelBlockable()) {
-                bar.setEnabled(false);
+                bar.setSwitchBarEnabled(false);
             }
 
             if (mChannelGroup != null && !isChannelGroupBlockable()) {
-                bar.setEnabled(false);
+                bar.setSwitchBarEnabled(false);
             }
 
             if (mChannel == null && mAppRow.systemApp
                     && (!mAppRow.banned || mAppRow.lockedImportance)) {
-                bar.setEnabled(false);
+                bar.setSwitchBarEnabled(false);
             }
 
             if (mChannel != null) {
diff --git a/src/com/android/settings/notification/zen/ZenModeButtonPreferenceController.java b/src/com/android/settings/notification/zen/ZenModeButtonPreferenceController.java
index dc338bd5ac2..a4318835786 100644
--- a/src/com/android/settings/notification/zen/ZenModeButtonPreferenceController.java
+++ b/src/com/android/settings/notification/zen/ZenModeButtonPreferenceController.java
@@ -16,37 +16,33 @@
 
 package com.android.settings.notification.zen;
 
-import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_FOCUSED;
-
 import android.app.settings.SettingsEnums;
 import android.content.Context;
 import android.provider.Settings;
-import android.view.View;
-import android.widget.Button;
+import android.widget.Switch;
 
 import androidx.fragment.app.FragmentManager;
 import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
 
-import com.android.settings.R;
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settings.dashboard.DashboardFragment;
 import com.android.settings.notification.SettingsEnableZenModeDialog;
 import com.android.settingslib.core.lifecycle.Lifecycle;
-import com.android.settingslib.widget.LayoutPreference;
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
 
 public class ZenModeButtonPreferenceController extends AbstractZenModePreferenceController
-        implements PreferenceControllerMixin {
+        implements PreferenceControllerMixin, OnMainSwitchChangeListener {
+
+    private static final String TAG = "EnableZenModeButton";
 
     public static final String KEY = "zen_mode_toggle";
 
-    private static final String TAG = "EnableZenModeButton";
     private final FragmentManager mFragment;
 
-    // DND can also be toggled from QS. If DND wasn't toggled by this preference, don't
-    // reroute focus.
-    private boolean mRefocusButton = false;
-    private Button mZenButtonOn;
-    private Button mZenButtonOff;
+    // DND can also be toggled from QS.
+    private MainSwitchPreference mPreference;
 
     public ZenModeButtonPreferenceController(Context context, Lifecycle lifecycle, FragmentManager
             fragment) {
@@ -64,26 +60,27 @@ public class ZenModeButtonPreferenceController extends AbstractZenModePreference
         return KEY;
     }
 
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        mPreference = (MainSwitchPreference) screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
+    }
+
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        if (isChecked) {
+            updateZenModeState(mPreference);
+        } else {
+            writeMetrics(mPreference, false);
+            mBackend.setZenMode(Settings.Global.ZEN_MODE_OFF);
+        }
+    }
+
     @Override
     public void updateState(Preference preference) {
         super.updateState(preference);
 
-        if (null == mZenButtonOn) {
-            mZenButtonOn = ((LayoutPreference) preference)
-                    .findViewById(R.id.zen_mode_settings_turn_on_button);
-            updateZenButtonOnClickListener(preference);
-        }
-
-        if (null == mZenButtonOff) {
-            mZenButtonOff = ((LayoutPreference) preference)
-                    .findViewById(R.id.zen_mode_settings_turn_off_button);
-            mZenButtonOff.setOnClickListener(v -> {
-                mRefocusButton = true;
-                writeMetrics(preference, false);
-                mBackend.setZenMode(Settings.Global.ZEN_MODE_OFF);
-            });
-        }
-
         updatePreference(preference);
     }
 
@@ -92,41 +89,27 @@ public class ZenModeButtonPreferenceController extends AbstractZenModePreference
             case Settings.Global.ZEN_MODE_ALARMS:
             case Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS:
             case Settings.Global.ZEN_MODE_NO_INTERRUPTIONS:
-                mZenButtonOff.setVisibility(View.VISIBLE);
-                mZenButtonOn.setVisibility(View.GONE);
-                if (mRefocusButton) {
-                    mRefocusButton = false;
-                    mZenButtonOff.sendAccessibilityEvent(TYPE_VIEW_FOCUSED);
-                }
+                mPreference.updateStatus(true);
                 break;
             case Settings.Global.ZEN_MODE_OFF:
             default:
-                mZenButtonOff.setVisibility(View.GONE);
-                updateZenButtonOnClickListener(preference);
-                mZenButtonOn.setVisibility(View.VISIBLE);
-                if (mRefocusButton) {
-                    mRefocusButton = false;
-                    mZenButtonOn.sendAccessibilityEvent(TYPE_VIEW_FOCUSED);
-                }
+                mPreference.updateStatus(false);
         }
     }
 
-    private void updateZenButtonOnClickListener(Preference preference) {
-        mZenButtonOn.setOnClickListener(v -> {
-            mRefocusButton = true;
-            writeMetrics(preference, true);
-            int zenDuration = getZenDuration();
-            switch (zenDuration) {
-                case Settings.Secure.ZEN_DURATION_PROMPT:
-                    new SettingsEnableZenModeDialog().show(mFragment, TAG);
-                    break;
-                case Settings.Secure.ZEN_DURATION_FOREVER:
-                    mBackend.setZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS);
-                    break;
-                default:
-                    mBackend.setZenModeForDuration(zenDuration);
-            }
-        });
+    private void updateZenModeState(Preference preference) {
+        writeMetrics(preference, true);
+        int zenDuration = getZenDuration();
+        switch (zenDuration) {
+            case Settings.Secure.ZEN_DURATION_PROMPT:
+                new SettingsEnableZenModeDialog().show(mFragment, TAG);
+                break;
+            case Settings.Secure.ZEN_DURATION_FOREVER:
+                mBackend.setZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS);
+                break;
+            default:
+                mBackend.setZenModeForDuration(zenDuration);
+        }
     }
 
     private void writeMetrics(Preference preference, boolean buttonOn) {
@@ -135,4 +118,4 @@ public class ZenModeButtonPreferenceController extends AbstractZenModePreference
         mMetricsFeatureProvider.action(mContext, SettingsEnums.ACTION_ZEN_TOGGLE_DND_BUTTON,
                 buttonOn);
     }
-}
\ No newline at end of file
+}
diff --git a/src/com/android/settings/widget/SettingsMainSwitchPreference.java b/src/com/android/settings/widget/SettingsMainSwitchPreference.java
index 4597842f6e3..2a7aa7384a6 100644
--- a/src/com/android/settings/widget/SettingsMainSwitchPreference.java
+++ b/src/com/android/settings/widget/SettingsMainSwitchPreference.java
@@ -17,8 +17,11 @@
 package com.android.settings.widget;
 
 import android.content.Context;
+import android.content.res.TypedArray;
+import android.text.TextUtils;
 import android.util.AttributeSet;
 
+import androidx.core.content.res.TypedArrayUtils;
 import androidx.preference.PreferenceViewHolder;
 import androidx.preference.TwoStatePreference;
 
@@ -47,23 +50,23 @@ public class SettingsMainSwitchPreference extends TwoStatePreference {
 
     public SettingsMainSwitchPreference(Context context) {
         super(context);
-        init();
+        init(context, null);
     }
 
     public SettingsMainSwitchPreference(Context context, AttributeSet attrs) {
         super(context, attrs);
-        init();
+        init(context, attrs);
     }
 
     public SettingsMainSwitchPreference(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
-        init();
+        init(context, attrs);
     }
 
     public SettingsMainSwitchPreference(Context context, AttributeSet attrs, int defStyleAttr,
             int defStyleRes) {
         super(context, attrs, defStyleAttr, defStyleRes);
-        init();
+        init(context, attrs);
     }
 
     @Override
@@ -79,8 +82,21 @@ public class SettingsMainSwitchPreference extends TwoStatePreference {
         registerListenerToSwitchBar();
     }
 
-    private void init() {
+    private void init(Context context, AttributeSet attrs) {
         setLayoutResource(R.layout.preference_widget_main_switch);
+
+        if (attrs != null) {
+            TypedArray a = context.obtainStyledAttributes(attrs,
+                    androidx.preference.R.styleable.Preference, 0/*defStyleAttr*/,
+                    0/*defStyleRes*/);
+            final CharSequence title = TypedArrayUtils.getText(a,
+                    androidx.preference.R.styleable.Preference_title,
+                    androidx.preference.R.styleable.Preference_android_title);
+            if (!TextUtils.isEmpty(title)) {
+                setTitle(title.toString());
+            }
+            a.recycle();
+        }
     }
 
     /**
@@ -153,7 +169,7 @@ public class SettingsMainSwitchPreference extends TwoStatePreference {
     /**
      * Enable or disable the text and switch.
      */
-    public void setEnabled(boolean enabled) {
+    public void setSwitchBarEnabled(boolean enabled) {
         if (mMainSwitchBar != null) {
             mMainSwitchBar.setEnabled(enabled);
         }
diff --git a/tests/robotests/src/com/android/settings/display/NightDisplayActivationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/NightDisplayActivationPreferenceControllerTest.java
index bad112a4388..d084c512938 100644
--- a/tests/robotests/src/com/android/settings/display/NightDisplayActivationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/NightDisplayActivationPreferenceControllerTest.java
@@ -21,13 +21,11 @@ import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.hardware.display.ColorDisplayManager;
-import android.view.View;
 
 import androidx.preference.PreferenceScreen;
 
-import com.android.settings.R;
 import com.android.settings.testutils.shadow.SettingsShadowResources;
-import com.android.settingslib.widget.LayoutPreference;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.After;
 import org.junit.Before;
@@ -45,7 +43,7 @@ public class NightDisplayActivationPreferenceControllerTest {
 
     @Mock
     private PreferenceScreen mScreen;
-    private LayoutPreference mPreference;
+    private MainSwitchPreference mPreference;
     private Context mContext;
     private ColorDisplayManager mColorDisplayManager;
     private NightDisplayActivationPreferenceController mPreferenceController;
@@ -55,7 +53,7 @@ public class NightDisplayActivationPreferenceControllerTest {
         MockitoAnnotations.initMocks(this);
         mContext = RuntimeEnvironment.application;
         mColorDisplayManager = mContext.getSystemService(ColorDisplayManager.class);
-        mPreference = new LayoutPreference(mContext, R.layout.night_display_activation_button);
+        mPreference = new MainSwitchPreference(mContext);
         when(mScreen.findPreference(anyString())).thenReturn(mPreference);
         mPreferenceController = new NightDisplayActivationPreferenceController(mContext,
                 "night_display_activation");
@@ -104,9 +102,9 @@ public class NightDisplayActivationPreferenceControllerTest {
     public void onClick_activates() {
         mColorDisplayManager.setNightDisplayActivated(false);
 
-        final View view = mPreference.findViewById(R.id.night_display_turn_on_button);
-        assertThat(view.getVisibility()).isEqualTo(View.VISIBLE);
-        view.performClick();
+        final NightDisplayActivationPreferenceController controller =
+                new NightDisplayActivationPreferenceController(mContext, "night_display_activated");
+        controller.onSwitchChanged(null, true);
 
         assertThat(mColorDisplayManager.isNightDisplayActivated()).isEqualTo(true);
     }
@@ -115,9 +113,9 @@ public class NightDisplayActivationPreferenceControllerTest {
     public void onClick_deactivates() {
         mColorDisplayManager.setNightDisplayActivated(true);
 
-        final View view = mPreference.findViewById(R.id.night_display_turn_off_button);
-        assertThat(view.getVisibility()).isEqualTo(View.VISIBLE);
-        view.performClick();
+        final NightDisplayActivationPreferenceController controller =
+                new NightDisplayActivationPreferenceController(mContext, "night_display_activated");
+        controller.onSwitchChanged(null, false);
 
         assertThat(mColorDisplayManager.isNightDisplayActivated()).isEqualTo(false);
     }
diff --git a/tests/robotests/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceControllerTest.java
index 33d0927c5d3..0d4646e9106 100644
--- a/tests/robotests/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceControllerTest.java
@@ -23,9 +23,8 @@ import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
 import android.app.UiModeManager;
@@ -33,13 +32,12 @@ import android.content.Context;
 import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.os.PowerManager;
-import android.view.View;
-import android.widget.Button;
 
 import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
-import com.android.settingslib.widget.LayoutPreference;
+import com.android.settings.testutils.FakeFeatureFactory;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -47,6 +45,7 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 import java.util.Locale;
 
@@ -54,25 +53,21 @@ import java.util.Locale;
 public class DarkModeActivationPreferenceControllerTest {
     private DarkModeActivationPreferenceController mController;
     private String mPreferenceKey = "key";
+
     @Mock
-    private LayoutPreference mPreference;
+    private MainSwitchPreference mPreference;
     @Mock
     private PreferenceScreen mScreen;
     @Mock
     private Resources res;
     @Mock
-    private Context mContext;
-    @Mock
     private UiModeManager mService;
     @Mock
-    private Button mTurnOffButton;
-    @Mock
-    private Button mTurnOnButton;
-    @Mock
     private PowerManager mPM;
     @Mock
     private TimeFormatter mFormat;
 
+    private Context mContext;
     private Configuration mConfigNightYes = new Configuration();
     private Configuration mConfigNightNo = new Configuration();
     private Locale mLocal = new Locale("ENG");
@@ -80,16 +75,14 @@ public class DarkModeActivationPreferenceControllerTest {
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+        FakeFeatureFactory.setupForTest();
+        mContext = spy(RuntimeEnvironment.application);
         mService = mock(UiModeManager.class);
         when(mContext.getResources()).thenReturn(res);
         when(res.getConfiguration()).thenReturn(mConfigNightNo);
         when(mContext.getSystemService(UiModeManager.class)).thenReturn(mService);
         when(mContext.getSystemService(PowerManager.class)).thenReturn(mPM);
         when(mScreen.findPreference(anyString())).thenReturn(mPreference);
-        when(mPreference.findViewById(
-                eq(R.id.dark_ui_turn_on_button))).thenReturn(mTurnOnButton);
-        when(mPreference.findViewById(
-                eq(R.id.dark_ui_turn_off_button))).thenReturn(mTurnOffButton);
         when(mService.setNightModeActivated(anyBoolean())).thenReturn(true);
         when(mFormat.of(any())).thenReturn("10:00 AM");
         when(mContext.getString(
@@ -124,66 +117,55 @@ public class DarkModeActivationPreferenceControllerTest {
     public void nightMode_toggleButton_offManual() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_YES);
         when(res.getConfiguration()).thenReturn(mConfigNightYes);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
 
         mController.updateState(mPreference);
 
-        verify(mTurnOnButton).setVisibility(eq(View.GONE));
-        verify(mTurnOffButton).setVisibility(eq(View.VISIBLE));
-        verify(mTurnOffButton).setText(eq(mContext.getString(
-                R.string.dark_ui_activation_off_manual)));
+        assertThat(preference.isChecked()).isFalse();
     }
 
     @Test
     public void nightMode_toggleButton_offCustom() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
         when(res.getConfiguration()).thenReturn(mConfigNightYes);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
 
-        mController.updateState(mPreference);
+        mController.updateState(preference);
 
-        verify(mTurnOnButton).setVisibility(eq(View.GONE));
-        verify(mTurnOffButton).setVisibility(eq(View.VISIBLE));
-        verify(mTurnOffButton).setText(eq(mContext.getString(
-                R.string.dark_ui_activation_off_custom)));
+        assertThat(preference.isChecked()).isFalse();
     }
 
     @Test
     public void nightMode_toggleButton_onCustom() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
         when(res.getConfiguration()).thenReturn(mConfigNightYes);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
 
-        mController.updateState(mPreference);
+        mController.updateState(preference);
 
-        verify(mTurnOnButton).setVisibility(eq(View.GONE));
-        verify(mTurnOffButton).setVisibility(eq(View.VISIBLE));
-        verify(mTurnOffButton).setText(eq(mContext.getString(
-                R.string.dark_ui_activation_on_custom)));
+        assertThat(preference.isChecked()).isFalse();
     }
 
-
     @Test
     public void nightMode_toggleButton_onAutoWhenModeIsYes() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_YES);
         when(res.getConfiguration()).thenReturn(mConfigNightNo);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext, null);
 
-        mController.updateState(mPreference);
+        mController.updateState(preference);
 
-        verify(mTurnOffButton).setVisibility(eq(View.GONE));
-        verify(mTurnOnButton).setVisibility(eq(View.VISIBLE));
-        verify(mTurnOnButton).setText(eq(mContext.getString(
-                R.string.dark_ui_activation_on_manual)));
+        assertThat(preference.isChecked()).isFalse();
     }
 
     @Test
     public void nightMode_toggleButton_onAutoWhenModeIsAuto() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_AUTO);
         when(res.getConfiguration()).thenReturn(mConfigNightNo);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
 
-        mController.updateState(mPreference);
+        mController.updateState(preference);
 
-        verify(mTurnOffButton).setVisibility(eq(View.GONE));
-        verify(mTurnOnButton).setVisibility(eq(View.VISIBLE));
-        verify(mTurnOnButton).setText(eq(mContext.getString(
-                R.string.dark_ui_activation_on_auto)));
+        assertThat(preference.isChecked()).isFalse();
     }
 
     @Test
@@ -207,9 +189,10 @@ public class DarkModeActivationPreferenceControllerTest {
     @Test
     public void buttonVisisbility_hideButton_offWhenInPowerSaveMode() {
         when(mPM.isPowerSaveMode()).thenReturn(true);
-        mController.updateState(mPreference);
-        verify(mTurnOffButton).setVisibility(eq(View.GONE));
-        verify(mTurnOnButton).setVisibility(eq(View.GONE));
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
+
+        mController.updateState(preference);
+        assertThat(preference.isChecked()).isFalse();
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/SmartBatteryPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/SmartBatteryPreferenceControllerTest.java
index 91e7b0486d3..a1f3bda371c 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/SmartBatteryPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/SmartBatteryPreferenceControllerTest.java
@@ -21,12 +21,12 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.doReturn;
 
 import android.content.ContentResolver;
+import android.content.Context;
 import android.provider.Settings;
 
-import androidx.preference.SwitchPreference;
-
 import com.android.settings.core.BasePreferenceController;
 import com.android.settings.testutils.FakeFeatureFactory;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -42,48 +42,50 @@ public class SmartBatteryPreferenceControllerTest {
     private static final int OFF = 0;
 
     private SmartBatteryPreferenceController mController;
-    private SwitchPreference mPreference;
     private ContentResolver mContentResolver;
     private FakeFeatureFactory mFeatureFactory;
+    private Context mContext;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
+        mContext = RuntimeEnvironment.application;
         mFeatureFactory = FakeFeatureFactory.setupForTest();
         mContentResolver = RuntimeEnvironment.application.getContentResolver();
         mController = new SmartBatteryPreferenceController(RuntimeEnvironment.application);
-        mPreference = new SwitchPreference(RuntimeEnvironment.application);
     }
 
     @Test
     public void testUpdateState_smartBatteryOn_preferenceChecked() {
         putSmartBatteryValue(ON);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
 
-        mController.updateState(mPreference);
+        mController.updateState(preference);
 
-        assertThat(mPreference.isChecked()).isTrue();
+        assertThat(preference.isChecked()).isTrue();
     }
 
     @Test
     public void testUpdateState_smartBatteryOff_preferenceUnchecked() {
         putSmartBatteryValue(OFF);
+        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
 
-        mController.updateState(mPreference);
+        mController.updateState(preference);
 
-        assertThat(mPreference.isChecked()).isFalse();
+        assertThat(preference.isChecked()).isFalse();
     }
 
     @Test
     public void testUpdateState_checkPreference_smartBatteryOn() {
-        mController.onPreferenceChange(mPreference, true);
+        mController.onSwitchChanged(null, true);
 
         assertThat(getSmartBatteryValue()).isEqualTo(ON);
     }
 
     @Test
     public void testUpdateState_unCheckPreference_smartBatteryOff() {
-        mController.onPreferenceChange(mPreference, false);
+        mController.onSwitchChanged(null, false);
 
         assertThat(getSmartBatteryValue()).isEqualTo(OFF);
     }
diff --git a/tests/robotests/src/com/android/settings/nfc/NfcEnablerTest.java b/tests/robotests/src/com/android/settings/nfc/NfcEnablerTest.java
index c3dfe48a775..cc451906214 100644
--- a/tests/robotests/src/com/android/settings/nfc/NfcEnablerTest.java
+++ b/tests/robotests/src/com/android/settings/nfc/NfcEnablerTest.java
@@ -26,7 +26,7 @@ import android.content.Context;
 import android.nfc.NfcAdapter;
 import android.provider.Settings;
 
-import androidx.preference.SwitchPreference;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -40,7 +40,7 @@ import org.robolectric.RuntimeEnvironment;
 public class NfcEnablerTest {
 
     @Mock
-    private SwitchPreference mNfcPreference;
+    private MainSwitchPreference mNfcPreference;
 
     private Context mContext;
     private NfcEnabler mNfcEnabler;
diff --git a/tests/robotests/src/com/android/settings/nfc/NfcPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/nfc/NfcPreferenceControllerTest.java
index 1d0516230a4..3ca553bc3ba 100644
--- a/tests/robotests/src/com/android/settings/nfc/NfcPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/nfc/NfcPreferenceControllerTest.java
@@ -31,11 +31,11 @@ import android.os.UserManager;
 import android.provider.Settings;
 
 import androidx.preference.PreferenceScreen;
-import androidx.preference.SwitchPreference;
 
 import com.android.settings.nfc.NfcPreferenceController.NfcSliceWorker;
 import com.android.settings.nfc.NfcPreferenceController.NfcSliceWorker.NfcUpdateReceiver;
 import com.android.settings.testutils.shadow.ShadowNfcAdapter;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -63,7 +63,7 @@ public class NfcPreferenceControllerTest {
     private PreferenceScreen mScreen;
 
     private Context mContext;
-    private SwitchPreference mNfcPreference;
+    private MainSwitchPreference mNfcPreference;
     private NfcPreferenceController mNfcController;
     private ShadowNfcAdapter mShadowNfcAdapter;
     private NfcAdapter mNfcAdapter;
@@ -81,7 +81,7 @@ public class NfcPreferenceControllerTest {
 
         mNfcController = new NfcPreferenceController(mContext,
                 NfcPreferenceController.KEY_TOGGLE_NFC);
-        mNfcPreference = new SwitchPreference(RuntimeEnvironment.application);
+        mNfcPreference = new MainSwitchPreference(RuntimeEnvironment.application);
 
         when(mScreen.findPreference(mNfcController.getPreferenceKey())).thenReturn(mNfcPreference);
     }
diff --git a/tests/robotests/src/com/android/settings/notification/zen/ZenModeButtonPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/zen/ZenModeButtonPreferenceControllerTest.java
index 12004026078..9daef1aeeb2 100644
--- a/tests/robotests/src/com/android/settings/notification/zen/ZenModeButtonPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/zen/ZenModeButtonPreferenceControllerTest.java
@@ -22,24 +22,21 @@ import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
 import static android.provider.Settings.Global.ZEN_MODE_NO_INTERRUPTIONS;
 import static android.provider.Settings.Global.ZEN_MODE_OFF;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.NotificationManager;
 import android.content.ContentResolver;
 import android.content.Context;
 import android.provider.Settings;
-import android.view.View;
-import android.widget.Button;
 
 import androidx.fragment.app.FragmentManager;
-import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
 
-import com.android.settings.notification.zen.ZenModeBackend;
-import com.android.settings.notification.zen.ZenModeButtonPreferenceController;
 import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -61,14 +58,10 @@ public class ZenModeButtonPreferenceControllerTest {
     @Mock
     private NotificationManager mNotificationManager;
     @Mock
-    private Preference mockPref;
+    private MainSwitchPreference mMockPref;
     @Mock
     private NotificationManager.Policy mPolicy;
     @Mock
-    private Button mZenButtonOn;
-    @Mock
-    private Button mZenButtonOff;
-    @Mock
     private PreferenceScreen mPreferenceScreen;
     private ContentResolver mContentResolver;
     private Context mContext;
@@ -85,67 +78,65 @@ public class ZenModeButtonPreferenceControllerTest {
                 mock(FragmentManager.class));
         when(mNotificationManager.getNotificationPolicy()).thenReturn(mPolicy);
         ReflectionHelpers.setField(mController, "mBackend", mBackend);
-        ReflectionHelpers.setField(mController, "mZenButtonOn", mZenButtonOn);
-        ReflectionHelpers.setField(mController, "mZenButtonOff", mZenButtonOff);
 
-        when(mPreferenceScreen.findPreference(mController.getPreferenceKey())).thenReturn(mockPref);
+        when(mPreferenceScreen.findPreference(mController.getPreferenceKey())).thenReturn(
+                mMockPref);
         mController.displayPreference(mPreferenceScreen);
     }
 
     @Test
     public void updateState_TotalSilence() {
         Settings.Global.putInt(mContentResolver, ZEN_MODE, ZEN_MODE_NO_INTERRUPTIONS);
-        final Preference mockPref = mock(Preference.class);
-        mController.updateState(mockPref);
+        final MainSwitchPreference pref = new MainSwitchPreference(mContext);
 
-        verify(mZenButtonOn).setVisibility(View.GONE);
-        verify(mZenButtonOff).setVisibility(View.VISIBLE);
+        mController.updateState(pref);
+
+        assertThat(pref.isChecked()).isFalse();
     }
 
     @Test
     public void updateState_AlarmsOnly() {
         Settings.Global.putInt(mContentResolver, ZEN_MODE, ZEN_MODE_ALARMS);
-        final Preference mockPref = mock(Preference.class);
-        mController.updateState(mockPref);
+        final MainSwitchPreference pref = new MainSwitchPreference(mContext);
 
-        verify(mZenButtonOn).setVisibility(View.GONE);
-        verify(mZenButtonOff).setVisibility(View.VISIBLE);
+        mController.updateState(pref);
+
+        assertThat(pref.isChecked()).isFalse();
     }
 
     @Test
     public void updateState_Priority() {
         Settings.Global.putInt(mContentResolver, ZEN_MODE, ZEN_MODE_IMPORTANT_INTERRUPTIONS);
-        final Preference mockPref = mock(Preference.class);
-        mController.updateState(mockPref);
+        final MainSwitchPreference pref = new MainSwitchPreference(mContext);
 
-        verify(mZenButtonOn).setVisibility(View.GONE);
-        verify(mZenButtonOff).setVisibility(View.VISIBLE);
+        mController.updateState(pref);
+
+        assertThat(pref.isChecked()).isFalse();
     }
 
     @Test
     public void updateState_ZenOff() {
         Settings.Global.putInt(mContentResolver, ZEN_MODE, ZEN_MODE_OFF);
-        final Preference mockPref = mock(Preference.class);
-        mController.updateState(mockPref);
+        final MainSwitchPreference pref = new MainSwitchPreference(mContext);
 
-        verify(mZenButtonOn).setVisibility(View.VISIBLE);
-        verify(mZenButtonOff).setVisibility(View.GONE);
+        mController.updateState(pref);
+
+        assertThat(pref.isChecked()).isFalse();
     }
 
     @Test
     public void updateState_otherUserChangedZen() {
-        final Preference mockPref = mock(Preference.class);
         Settings.Global.putInt(mContentResolver, ZEN_MODE, ZEN_MODE_OFF);
-        mController.updateState(mockPref);
-        verify(mZenButtonOn).setVisibility(View.VISIBLE);
-        verify(mZenButtonOff).setVisibility(View.GONE);
+        final MainSwitchPreference pref = new MainSwitchPreference(mContext);
+        mController.updateState(pref);
+
+        assertThat(pref.isChecked()).isFalse();
 
         Settings.Global.putInt(mContentResolver, ZEN_MODE, ZEN_MODE_IMPORTANT_INTERRUPTIONS);
         final int GUEST_USER_ID = 10;
         mController.mSettingObserver.onChange(false,
                 Settings.Global.getUriFor(Settings.Global.ZEN_MODE), GUEST_USER_ID);
 
-        verify(mZenButtonOn).setVisibility(View.GONE);
-        verify(mZenButtonOff).setVisibility(View.VISIBLE);
+        assertThat(pref.isChecked()).isFalse();
     }
-}
\ No newline at end of file
+}