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
diff --git a/res/layout/dark_ui_activation_button.xml b/res/layout/dark_ui_activation_button.xml
deleted file mode 100644
index 5f9eefc..0000000
--- 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 df160ec..0000000
--- 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 db889ff..0000000
--- 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 09aea0e..8a811c0 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 6b304ef..ebcf174 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 19279e0..5ac9e26 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 8ff4983..09ce277 100644
--- a/res/xml/nfc_and_payment_settings.xml
+++ b/res/xml/nfc_and_payment_settings.xml
@@ -19,6 +19,11 @@
     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"
@@ -26,12 +31,6 @@
         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"
         settings:controller="com.android.settings.nfc.SecureNfcPreferenceController"
diff --git a/res/xml/night_display_settings.xml b/res/xml/night_display_settings.xml
index 97f5dd2..23afccf 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 87eeb86..2d36119 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"
@@ -30,13 +35,6 @@
         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"
         android:summary="@string/battery_auto_restriction_summary"
@@ -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 6764e28..78dee02 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 030c5b9..3138ae0 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 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 cf9332b..e6217ff 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 @@
 
     @Override
     public int getAvailabilityStatus() {
-        return ColorDisplayManager.isNightDisplayAvailable(mContext) ? AVAILABLE
+        return ColorDisplayManager.isNightDisplayAvailable(mContext) ? AVAILABLE_UNSEARCHABLE
                 : UNSUPPORTED_ON_DEVICE;
     }
 
@@ -82,16 +69,24 @@
     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 @@
         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 1027dac..800e7e5 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.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.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 @@
 
     @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 @@
         }
     }
 
-    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 040af9f..ea277a1 100644
--- a/src/com/android/settings/fuelgauge/SmartBatteryPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/SmartBatteryPreferenceController.java
@@ -20,22 +20,27 @@
 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 @@
         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 777e7d1..88482d5 100644
--- a/src/com/android/settings/nfc/NfcEnabler.java
+++ b/src/com/android/settings/nfc/NfcEnabler.java
@@ -21,16 +21,17 @@
 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 @@
     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 9bf013c..b30e7c3f 100644
--- a/src/com/android/settings/nfc/NfcPreferenceController.java
+++ b/src/com/android/settings/nfc/NfcPreferenceController.java
@@ -23,24 +23,27 @@
 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 @@
             return;
         }
 
-        final SwitchPreference switchPreference = screen.findPreference(getPreferenceKey());
+        mPreference = (MainSwitchPreference) screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
+        mNfcEnabler = new NfcEnabler(mContext, mPreference);
+    }
 
-        mNfcEnabler = new NfcEnabler(mContext, switchPreference);
-
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        if (isChecked != mNfcAdapter.isEnabled()) {
+            setChecked(isChecked);
+        }
     }
 
     @Override
@@ -129,12 +138,12 @@
      * 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 @@
                         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 e5a7e98..4df3e28 100644
--- a/src/com/android/settings/notification/app/BlockPreferenceController.java
+++ b/src/com/android/settings/notification/app/BlockPreferenceController.java
@@ -72,16 +72,16 @@
             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 dc338bd..a431883 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) {
@@ -65,25 +61,26 @@
     }
 
     @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 @@
             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 @@
         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 4597842..2a7aa73 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 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 @@
         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 @@
     /**
      * 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 bad112a..d084c51 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 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 @@
 
     @Mock
     private PreferenceScreen mScreen;
-    private LayoutPreference mPreference;
+    private MainSwitchPreference mPreference;
     private Context mContext;
     private ColorDisplayManager mColorDisplayManager;
     private NightDisplayActivationPreferenceController mPreferenceController;
@@ -55,7 +53,7 @@
         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 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 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 33d0927c..0d4646e 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.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.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.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 import java.util.Locale;
 
@@ -54,25 +53,21 @@
 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 @@
     @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 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 @@
     @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 91e7b04..a1f3bda 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 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 @@
     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 c3dfe48..cc45190 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.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 @@
 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 1d05162..3ca553b 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.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 @@
     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 @@
 
         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 1200402..9daef1a 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_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 @@
     @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 @@
                 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
+}