From 9b906cb5b6e13bfad9153c36d229877795e2fe94 Mon Sep 17 00:00:00 2001
From: Stanley Wang <stanleytfwang@google.com>
Date: Mon, 12 Apr 2021 19:46:22 +0800
Subject: [PATCH] Use MainSwitchPreference on Bubbles, Screen Saver and
 One-Handed mode pages.

Fixes: 178679876
Fixes: 177968619
Test: robotest and see the UIs.
Change-Id: Ic62c7464a71e5410ece5d1db7578c522e1babedc
---
 res/layout/dream_start_button.xml             | 38 ------------
 res/values/strings.xml                        | 10 ++-
 res/xml/bubble_notification_settings.xml      | 16 ++---
 res/xml/dream_fragment_overview.xml           | 15 ++---
 res/xml/one_handed_settings.xml               | 12 ++--
 .../android/settings/dream/DreamSettings.java |  3 +-
 .../dream/StartNowPreferenceController.java   | 61 ++++++++-----------
 .../OneHandedEnablePreferenceController.java  |  4 +-
 ...ubbleNotificationPreferenceController.java | 30 +++++----
 ...ettingsMainSwitchPreferenceController.java | 56 +++++++++++++++++
 .../StartNowPreferenceControllerTest.java     | 50 +++++++--------
 ...eNotificationPreferenceControllerTest.java | 60 ++++++++----------
 12 files changed, 169 insertions(+), 186 deletions(-)
 delete mode 100644 res/layout/dream_start_button.xml
 create mode 100644 src/com/android/settings/widget/SettingsMainSwitchPreferenceController.java

diff --git a/res/layout/dream_start_button.xml b/res/layout/dream_start_button.xml
deleted file mode 100644
index a45e56ffa13..00000000000
--- a/res/layout/dream_start_button.xml
+++ /dev/null
@@ -1,38 +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:orientation="horizontal"
-    android:gravity="bottom"
-    android:paddingTop="4dp"
-    android:paddingStart="72dp"
-    android:paddingEnd="72dp"
-    android:layout_width="match_parent"
-    android:layout_height="match_parent">
-
-    <Button
-        android:id="@+id/dream_start_now_button"
-        style="@style/ActionSecondaryButton"
-        android:layout_width="0dp"
-        android:layout_weight="1"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center"
-        android:text="@string/screensaver_settings_dream_start"
-        android:paddingEnd="8dp" />
-
-</LinearLayout>
\ No newline at end of file
diff --git a/res/values/strings.xml b/res/values/strings.xml
index cb921ae7395..591fc488a62 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -3121,8 +3121,6 @@
     <string name="screensaver_settings_when_to_dream">When to start</string>
     <!-- Dream settings screen, action label, current selected screen saver -->
     <string name="screensaver_settings_current">Current screen saver</string>
-    <!-- Dream settings screen, button label to start dreaming -->
-    <string name="screensaver_settings_dream_start">Start now</string>
     <!-- Dream settings screen, button label for settings for a specific screensaver -->
     <string name="screensaver_settings_button">Settings</string>
     <!-- Sound & display settings screen, setting option name to change whether the screen adjusts automatically based on lighting conditions -->
@@ -8747,10 +8745,6 @@
 
     <!-- Bubble feature settings: Description for toggle when bubbles are turned on in notification settings [CHAR LIMIT=100] -->
     <string name="notifications_bubble_setting_on_summary">On / Conversations can appear as floating icons</string>
-    <!-- Bubble feature settings: Title of the toggle to turn bubbles on and off [CHAR LIMIT=100] -->
-    <string name="notifications_bubble_setting_title">Allow apps to show bubbles</string>
-    <!-- Bubble feature settings: Description shown with the toggle to turn bubbles on and off [CHAR LIMIT=NONE]-->
-    <string name="notifications_bubble_setting_description">Some conversations will appear as floating icons on top of other apps</string>
 
     <!-- Bubble app settings: option allowing all conversations from an app to bubble if the app sends bubbles [CHAR LIMIT=60]-->
     <string name="bubble_app_setting_all">All conversations can bubble</string>
@@ -13236,6 +13230,10 @@
     <string name="adaptive_brightness_main_switch_title">Use adaptive brightness</string>
     <!-- Title for wifi calling main switch preferences. [CHAR LIMIT=50] -->
     <string name="wifi_calling_main_switch_title">Use Wi‑Fi calling</string>
+    <!-- Title for Bubbles main switch preferences. [CHAR LIMIT=50] -->
+    <string name="bubbles_main_switch_title">Use Bubbles</string>
+    <!-- Title for Screen saver main switch preferences. [CHAR LIMIT=50] -->
+    <string name="screen_saver_main_switch_title">Use screen saver</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/bubble_notification_settings.xml b/res/xml/bubble_notification_settings.xml
index 4827cb6df95..70492b4fd8b 100644
--- a/res/xml/bubble_notification_settings.xml
+++ b/res/xml/bubble_notification_settings.xml
@@ -18,6 +18,11 @@
                   xmlns:settings="http://schemas.android.com/apk/res-auto"
                   android:title="@string/bubbles_app_toggle_title">
 
+        <com.android.settingslib.widget.MainSwitchPreference
+            android:key="global_notification_bubbles"
+            android:title="@string/bubbles_main_switch_title"
+            settings:controller="com.android.settings.notification.BubbleNotificationPreferenceController"/>
+
         <com.android.settings.widget.VideoPreference
             android:key="bubbles_illustration"
             android:title="@string/summary_placeholder"
@@ -25,13 +30,4 @@
             settings:controller="com.android.settings.widget.VideoPreferenceController"
             android:persistent="false" />
 
-        <!-- Notification bubbles -->
-        <SwitchPreference
-            android:key="global_notification_bubbles"
-            android:icon="@drawable/ic_create_bubble"
-            settings:allowDividerAbove="true"
-            android:title="@string/notifications_bubble_setting_title"
-            android:summary="@string/notifications_bubble_setting_description"
-            settings:controller="com.android.settings.notification.BubbleNotificationPreferenceController"/>
-
-</PreferenceScreen>
\ No newline at end of file
+</PreferenceScreen>
diff --git a/res/xml/dream_fragment_overview.xml b/res/xml/dream_fragment_overview.xml
index 3cb49f9cced..6e18a979f3e 100644
--- a/res/xml/dream_fragment_overview.xml
+++ b/res/xml/dream_fragment_overview.xml
@@ -19,6 +19,11 @@
     xmlns:settings="http://schemas.android.com/apk/res-auto"
     android:title="@string/screensaver_settings_title">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="dream_start_now_button_container"
+        android:title="@string/screen_saver_main_switch_title"
+        settings:controller="com.android.settings.dream.StartNowPreferenceController"/>
+
     <com.android.settings.widget.GearPreference
         android:key="current_screensaver"
         android:title="@string/screensaver_settings_current"
@@ -30,14 +35,4 @@
         android:title="@string/screensaver_settings_when_to_dream"
         android:fragment="com.android.settings.dream.WhenToDreamPicker" />
 
-
-    <!-- Layout preference doesn't obey allowDividerAbove, so put it in a PreferenceCategory -->
-    <PreferenceCategory
-        android:key="dream_start_now_category">
-        <com.android.settingslib.widget.LayoutPreference
-            android:key="dream_start_now_button_container"
-            android:selectable="false"
-            android:layout="@layout/dream_start_button" />
-    </PreferenceCategory>
-
 </PreferenceScreen>
diff --git a/res/xml/one_handed_settings.xml b/res/xml/one_handed_settings.xml
index 5a71d48cd4e..b2c7de9788d 100644
--- a/res/xml/one_handed_settings.xml
+++ b/res/xml/one_handed_settings.xml
@@ -21,6 +21,12 @@
     android:persistent="false"
     android:title="@string/one_handed_title">
 
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:key="gesture_one_handed_mode_enabled"
+        android:title="@string/one_handed_mode_enabled"
+        settings:keywords="@string/keywords_one_handed"
+        settings:controller="com.android.settings.gestures.OneHandedEnablePreferenceController"/>
+
     <com.android.settingslib.widget.LayoutPreference
         android:key="one_handed_header"
         android:layout="@layout/one_handed_header"
@@ -29,12 +35,6 @@
         settings:allowDividerBelow="false"
         settings:searchable="false"/>
 
-    <SwitchPreference
-        android:key="gesture_one_handed_mode_enabled"
-        android:title="@string/one_handed_mode_enabled"
-        settings:keywords="@string/keywords_one_handed"
-        settings:controller="com.android.settings.gestures.OneHandedEnablePreferenceController"/>
-
     <SwitchPreference
         android:key="gesture_app_taps_to_exit"
         android:title="@string/one_handed_app_taps_to_exit"
diff --git a/src/com/android/settings/dream/DreamSettings.java b/src/com/android/settings/dream/DreamSettings.java
index a40cf88da99..83979e5da85 100644
--- a/src/com/android/settings/dream/DreamSettings.java
+++ b/src/com/android/settings/dream/DreamSettings.java
@@ -129,9 +129,8 @@ public class DreamSettings extends DashboardFragment {
     }
 
     private static List<AbstractPreferenceController> buildPreferenceControllers(Context context) {
-        List<AbstractPreferenceController> controllers = new ArrayList<>();
+        final List<AbstractPreferenceController> controllers = new ArrayList<>();
         controllers.add(new WhenToDreamPreferenceController(context));
-        controllers.add(new StartNowPreferenceController(context));
         return controllers;
     }
 
diff --git a/src/com/android/settings/dream/StartNowPreferenceController.java b/src/com/android/settings/dream/StartNowPreferenceController.java
index fb0581f343a..ed0a4dbeece 100644
--- a/src/com/android/settings/dream/StartNowPreferenceController.java
+++ b/src/com/android/settings/dream/StartNowPreferenceController.java
@@ -17,63 +17,52 @@
 package com.android.settings.dream;
 
 import android.content.Context;
-import android.widget.Button;
 
 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.overlay.FeatureFactory;
-import com.android.settingslib.core.AbstractPreferenceController;
+import com.android.settings.widget.SettingsMainSwitchPreferenceController;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 import com.android.settingslib.dream.DreamBackend;
-import com.android.settingslib.widget.LayoutPreference;
 
-public class StartNowPreferenceController extends AbstractPreferenceController implements
-        PreferenceControllerMixin {
+/**
+ * Controller that used to enable screen saver
+ */
+public class StartNowPreferenceController extends SettingsMainSwitchPreferenceController {
 
-    private static final String PREF_KEY = "dream_start_now_button_container";
     private final DreamBackend mBackend;
     private final MetricsFeatureProvider mMetricsFeatureProvider;
 
-    public StartNowPreferenceController(Context context) {
-        super(context);
-
+    public StartNowPreferenceController(Context context, String preferenceKey) {
+        super(context, preferenceKey);
         mBackend = DreamBackend.getInstance(context);
         mMetricsFeatureProvider = FeatureFactory.getFactory(context).getMetricsFeatureProvider();
     }
 
     @Override
-    public boolean isAvailable() {
-        return true;
-    }
-
-    @Override
-    public String getPreferenceKey() {
-        return PREF_KEY;
-    }
-
-    @Override
-    public void displayPreference(PreferenceScreen screen) {
-        super.displayPreference(screen);
-
-        LayoutPreference pref = screen.findPreference(getPreferenceKey());
-        Button startButton = pref.findViewById(R.id.dream_start_now_button);
-        startButton.setOnClickListener(v -> {
-            mMetricsFeatureProvider.logClickedPreference(pref,
-                    pref.getExtras().getInt(DashboardFragment.CATEGORY));
-            mBackend.startDreaming();
-        });
+    public int getAvailabilityStatus() {
+        return AVAILABLE;
     }
 
     @Override
     public void updateState(Preference preference) {
-        super.updateState(preference);
+        mSwitchPreference.setChecked(false);
+        mSwitchPreference.setEnabled(mBackend.getWhenToDreamSetting() != DreamBackend.NEVER);
+    }
 
-        Button startButton = ((LayoutPreference) preference)
-                .findViewById(R.id.dream_start_now_button);
-        startButton.setEnabled(mBackend.getWhenToDreamSetting() != DreamBackend.NEVER);
+    @Override
+    public boolean isChecked() {
+        return false;
+    }
+
+    @Override
+    public boolean setChecked(boolean isChecked) {
+        if (isChecked) {
+            mMetricsFeatureProvider.logClickedPreference(mSwitchPreference,
+                    mSwitchPreference.getExtras().getInt(DashboardFragment.CATEGORY));
+            mBackend.startDreaming();
+        }
+        return true;
     }
 }
diff --git a/src/com/android/settings/gestures/OneHandedEnablePreferenceController.java b/src/com/android/settings/gestures/OneHandedEnablePreferenceController.java
index adfa137232e..75018970449 100644
--- a/src/com/android/settings/gestures/OneHandedEnablePreferenceController.java
+++ b/src/com/android/settings/gestures/OneHandedEnablePreferenceController.java
@@ -20,12 +20,12 @@ import android.content.Context;
 
 import com.android.settings.R;
 import com.android.settings.core.BasePreferenceController;
-import com.android.settings.core.TogglePreferenceController;
+import com.android.settings.widget.SettingsMainSwitchPreferenceController;
 
 /**
  * The controller to handle one-handed mode enable or disable state.
  **/
-public class OneHandedEnablePreferenceController extends TogglePreferenceController {
+public class OneHandedEnablePreferenceController extends SettingsMainSwitchPreferenceController {
 
     public OneHandedEnablePreferenceController(Context context, String key) {
         super(context, key);
diff --git a/src/com/android/settings/notification/BubbleNotificationPreferenceController.java b/src/com/android/settings/notification/BubbleNotificationPreferenceController.java
index 67f0c7b34ab..37723b13283 100644
--- a/src/com/android/settings/notification/BubbleNotificationPreferenceController.java
+++ b/src/com/android/settings/notification/BubbleNotificationPreferenceController.java
@@ -30,8 +30,7 @@ import androidx.annotation.VisibleForTesting;
 import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
 
-import com.android.settings.core.PreferenceControllerMixin;
-import com.android.settings.core.TogglePreferenceController;
+import com.android.settings.widget.SettingsMainSwitchPreferenceController;
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
 import com.android.settingslib.core.lifecycle.events.OnPause;
 import com.android.settingslib.core.lifecycle.events.OnResume;
@@ -40,9 +39,8 @@ import com.android.settingslib.core.lifecycle.events.OnResume;
  * Feature level screen for bubbles, available through notification menu.
  * Allows user to turn bubbles on or off for the device.
  */
-public class BubbleNotificationPreferenceController extends TogglePreferenceController
-        implements PreferenceControllerMixin, Preference.OnPreferenceChangeListener,
-        LifecycleObserver, OnResume, OnPause {
+public class BubbleNotificationPreferenceController extends
+        SettingsMainSwitchPreferenceController implements LifecycleObserver, OnResume, OnPause {
 
     private static final String TAG = "BubbleNotifPrefContr";
 
@@ -60,9 +58,8 @@ public class BubbleNotificationPreferenceController extends TogglePreferenceCont
     @Override
     public void displayPreference(PreferenceScreen screen) {
         super.displayPreference(screen);
-        Preference preference = screen.findPreference(getPreferenceKey());
-        if (preference != null) {
-            mSettingObserver = new SettingObserver(preference);
+        if (mSwitchPreference != null) {
+            mSettingObserver = new SettingObserver(mSwitchPreference);
         }
     }
 
@@ -86,21 +83,22 @@ public class BubbleNotificationPreferenceController extends TogglePreferenceCont
         return am.isLowRamDevice() ? UNSUPPORTED_ON_DEVICE : AVAILABLE;
     }
 
+    @Override
+    public boolean isSliceable() {
+        return false;
+    }
+
     @Override
     public boolean isChecked() {
-        return Settings.Secure.getInt(mContext.getContentResolver(),
+        return Settings.Global.getInt(mContext.getContentResolver(),
                 NOTIFICATION_BUBBLES, ON) == ON;
     }
 
     @Override
     public boolean setChecked(boolean isChecked) {
-        return Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES,
-                isChecked ? ON : OFF);
-    }
-
-    @Override
-    public boolean isSliceable() {
-        return false;
+        Settings.Global.putInt(mContext.getContentResolver(),
+                NOTIFICATION_BUBBLES, isChecked ? ON : OFF);
+        return true;
     }
 
     class SettingObserver extends ContentObserver {
diff --git a/src/com/android/settings/widget/SettingsMainSwitchPreferenceController.java b/src/com/android/settings/widget/SettingsMainSwitchPreferenceController.java
new file mode 100644
index 00000000000..5fa9259ab9e
--- /dev/null
+++ b/src/com/android/settings/widget/SettingsMainSwitchPreferenceController.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.widget;
+
+import android.content.Context;
+import android.widget.Switch;
+
+import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
+
+import com.android.settings.core.TogglePreferenceController;
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
+
+/**
+ * Preference controller for MainSwitchPreference.
+ */
+public abstract class SettingsMainSwitchPreferenceController extends
+        TogglePreferenceController implements OnMainSwitchChangeListener {
+
+    protected MainSwitchPreference mSwitchPreference;
+
+    public SettingsMainSwitchPreferenceController(Context context, String preferenceKey) {
+        super(context, preferenceKey);
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        final Preference pref = screen.findPreference(getPreferenceKey());
+        if (pref != null && pref instanceof MainSwitchPreference) {
+            mSwitchPreference = (MainSwitchPreference) pref;
+            mSwitchPreference.addOnSwitchChangeListener(this);
+        }
+    }
+
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        mSwitchPreference.setChecked(isChecked);
+        setChecked(isChecked);
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/dream/StartNowPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/dream/StartNowPreferenceControllerTest.java
index d58d7e8d5f3..6c8a8f60cbd 100644
--- a/tests/robotests/src/com/android/settings/dream/StartNowPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/dream/StartNowPreferenceControllerTest.java
@@ -16,43 +16,37 @@
 
 package com.android.settings.dream;
 
-import static com.google.common.truth.Truth.assertThat;
-
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
-import android.view.View.OnClickListener;
-import android.widget.Button;
 
 import androidx.preference.PreferenceScreen;
 
-import com.android.settings.R;
 import com.android.settingslib.dream.DreamBackend;
-import com.android.settingslib.widget.LayoutPreference;
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Answers;
-import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.util.ReflectionHelpers;
 
 @RunWith(RobolectricTestRunner.class)
 public class StartNowPreferenceControllerTest {
 
     private StartNowPreferenceController mController;
-    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     private Context mContext;
+
     @Mock
     private PreferenceScreen mScreen;
     @Mock
-    private LayoutPreference mLayoutPref;
-    @Mock
-    private Button mButton;
+    private MainSwitchPreference mPref;
     @Mock
     private DreamBackend mBackend;
 
@@ -60,36 +54,38 @@ public class StartNowPreferenceControllerTest {
     public void setup() {
         MockitoAnnotations.initMocks(this);
 
-        mController = new StartNowPreferenceController(mContext);
-        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mLayoutPref);
-        when(mLayoutPref.findViewById(R.id.dream_start_now_button)).thenReturn(mButton);
+        mContext = spy(RuntimeEnvironment.application);
+        mController = new StartNowPreferenceController(mContext, "key");
+        mPref = mock(MainSwitchPreference.class);
+        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mPref);
 
         ReflectionHelpers.setField(mController, "mBackend", mBackend);
     }
 
     @Test
-    public void setsOnClickListenerForStartNow() {
-        ArgumentCaptor<OnClickListener> captor =
-                ArgumentCaptor.forClass(Button.OnClickListener.class);
-
+    public void displayPreference_shouldAddOnSwitchChangeListener() {
         mController.displayPreference(mScreen);
-        verify(mButton).setOnClickListener(captor.capture());
-        assertThat(captor.getValue()).isNotNull();
+
+        verify(mPref).addOnSwitchChangeListener(mController);
     }
 
     @Test
-    public void buttonIsDisabledWhenNeverDreaming() {
+    public void updateState_neverDreaming_preferenceShouldDidabled() {
         when(mBackend.getWhenToDreamSetting()).thenReturn(DreamBackend.NEVER);
+        mController.displayPreference(mScreen);
 
-        mController.updateState(mLayoutPref);
-        verify(mButton).setEnabled(false);
+        mController.updateState(mPref);
+
+        verify(mPref).setEnabled(false);
     }
 
     @Test
-    public void buttonIsEnabledWhenDreamIsAvailable() {
+    public void updateState_dreamIsAvailable_preferenceShouldEnabled() {
         when(mBackend.getWhenToDreamSetting()).thenReturn(DreamBackend.EITHER);
+        mController.displayPreference(mScreen);
 
-        mController.updateState(mLayoutPref);
-        verify(mButton).setEnabled(true);
+        mController.updateState(mPref);
+
+        verify(mPref).setEnabled(true);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/notification/BubbleNotificationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/BubbleNotificationPreferenceControllerTest.java
index 4f38c461e6f..8d643ad4613 100644
--- a/tests/robotests/src/com/android/settings/notification/BubbleNotificationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/BubbleNotificationPreferenceControllerTest.java
@@ -26,17 +26,15 @@ import static com.android.settings.notification.BadgingNotificationPreferenceCon
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager;
 import android.content.Context;
 import android.provider.Settings;
 
-import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
-import androidx.preference.TwoStatePreference;
+
+import com.android.settingslib.widget.MainSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -57,7 +55,7 @@ public class BubbleNotificationPreferenceControllerTest {
     private PreferenceScreen mScreen;
 
     private BubbleNotificationPreferenceController mController;
-    private Preference mPreference;
+    private MainSwitchPreference mPreference;
 
     private static final String KEY_NOTIFICATION_BUBBLES = "notification_bubbles";
 
@@ -67,15 +65,16 @@ public class BubbleNotificationPreferenceControllerTest {
         mContext = RuntimeEnvironment.application;
         mController = new BubbleNotificationPreferenceController(mContext,
                 KEY_NOTIFICATION_BUBBLES);
-        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference = new MainSwitchPreference(RuntimeEnvironment.application);
         mPreference.setKey(mController.getPreferenceKey());
         when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
+        mController.displayPreference(mScreen);
     }
 
     @Test
     public void isAvailable_lowRam_returnsUnsupported() {
         final ShadowActivityManager activityManager =
-               Shadow.extract(mContext.getSystemService(ActivityManager.class));
+                Shadow.extract(mContext.getSystemService(ActivityManager.class));
         activityManager.setIsLowRamDevice(true);
         assertEquals(UNSUPPORTED_ON_DEVICE, mController.getAvailabilityStatus());
     }
@@ -83,45 +82,40 @@ public class BubbleNotificationPreferenceControllerTest {
     @Test
     public void isAvailable_notLowRam_returnsAvailable() {
         final ShadowActivityManager activityManager =
-               Shadow.extract(mContext.getSystemService(ActivityManager.class));
+                Shadow.extract(mContext.getSystemService(ActivityManager.class));
         activityManager.setIsLowRamDevice(false);
         assertEquals(AVAILABLE, mController.getAvailabilityStatus());
     }
 
-    @Test
-    public void updateState_settingIsOn_preferenceSetChecked() {
-        final TwoStatePreference preference = mock(TwoStatePreference.class);
-        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, ON);
-
-        mController.updateState(preference);
-
-        verify(preference).setChecked(true);
-    }
-
     @Test
     public void updateState_settingIsOff_preferenceSetUnchecked() {
-        final TwoStatePreference preference = mock(TwoStatePreference.class);
-        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, OFF);
-        assertThat(Settings.Secure.getInt(mContext.getContentResolver(),
+        Settings.Global.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, OFF);
+        assertThat(Settings.Global.getInt(mContext.getContentResolver(),
                 NOTIFICATION_BUBBLES, ON)).isEqualTo(OFF);
 
-        mController.updateState(preference);
+        mPreference.updateStatus(false);
 
-        verify(preference).setChecked(false);
+        assertThat(mPreference.isChecked()).isFalse();
     }
 
     @Test
-    public void isChecked_settingIsOff_shouldReturnFalse() {
-        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, OFF);
+    public void onSwitchChanged_true_settingIsOff_flagShouldOn() {
+        Settings.Global.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, OFF);
 
-        assertThat(mController.isChecked()).isFalse();
+        mController.onSwitchChanged(null, true);
+
+        assertThat(Settings.Global.getInt(mContext.getContentResolver(),
+                NOTIFICATION_BUBBLES, OFF)).isEqualTo(ON);
     }
 
     @Test
-    public void isChecked_settingIsOn_shouldReturnTrue() {
-        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, ON);
+    public void onSwitchChanged_false_settingIsOn_flagShouldOff() {
+        Settings.Global.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, ON);
 
-        assertThat(mController.isChecked()).isTrue();
+        mController.onSwitchChanged(null, false);
+
+        assertThat(Settings.Global.getInt(mContext.getContentResolver(),
+                NOTIFICATION_BUBBLES, ON)).isEqualTo(OFF);
     }
 
     @Test
@@ -129,8 +123,8 @@ public class BubbleNotificationPreferenceControllerTest {
         Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, ON);
 
         mController.setChecked(false);
-        int updatedValue = Settings.Secure.getInt(mContext.getContentResolver(),
-                NOTIFICATION_BUBBLES, -1);
+        int updatedValue = Settings.Global.getInt(mContext.getContentResolver(),
+                NOTIFICATION_BUBBLES, ON);
 
         assertThat(updatedValue).isEqualTo(OFF);
     }
@@ -140,8 +134,8 @@ public class BubbleNotificationPreferenceControllerTest {
         Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BUBBLES, OFF);
 
         mController.setChecked(true);
-        int updatedValue = Settings.Secure.getInt(mContext.getContentResolver(),
-                NOTIFICATION_BUBBLES, -1);
+        int updatedValue = Settings.Global.getInt(mContext.getContentResolver(),
+                NOTIFICATION_BUBBLES, OFF);
 
         assertThat(updatedValue).isEqualTo(ON);
     }