diff --git a/res/xml/configure_notification_settings.xml b/res/xml/configure_notification_settings.xml
index 612f5e7623b..291d9a8da6d 100644
--- a/res/xml/configure_notification_settings.xml
+++ b/res/xml/configure_notification_settings.xml
@@ -28,7 +28,8 @@
     <!-- Notification badging -->
     <SwitchPreference
         android:key="notification_badging"
-        android:title="@string/notification_badging_title"/>
+        android:title="@string/notification_badging_title"
+        settings:controller="com.android.settings.notification.BadgingNotificationPreferenceController"/>
 
     <!-- Pulse notification light -->
     <SwitchPreference
diff --git a/src/com/android/settings/notification/BadgingNotificationPreferenceController.java b/src/com/android/settings/notification/BadgingNotificationPreferenceController.java
index 18efd43be12..37e3db0fb5f 100644
--- a/src/com/android/settings/notification/BadgingNotificationPreferenceController.java
+++ b/src/com/android/settings/notification/BadgingNotificationPreferenceController.java
@@ -23,11 +23,13 @@ import android.database.ContentObserver;
 import android.net.Uri;
 import android.os.Handler;
 import android.provider.Settings;
+import android.support.annotation.VisibleForTesting;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 import android.support.v7.preference.TwoStatePreference;
 
 import com.android.settings.core.PreferenceControllerMixin;
+import com.android.settings.core.TogglePreferenceController;
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settings.R;
 import com.android.settings.search.DatabaseIndexingUtils;
@@ -39,19 +41,20 @@ import com.android.settingslib.core.lifecycle.events.OnResume;
 
 import static android.provider.Settings.Secure.NOTIFICATION_BADGING;
 
-public class BadgingNotificationPreferenceController extends AbstractPreferenceController
+public class BadgingNotificationPreferenceController extends TogglePreferenceController
         implements PreferenceControllerMixin, Preference.OnPreferenceChangeListener,
         LifecycleObserver, OnResume, OnPause {
 
     private static final String TAG = "BadgeNotifPrefContr";
-    private static final String KEY_NOTIFICATION_BADGING = "notification_badging";
-    private static final int ON = 1;
-    private static final int OFF = 0;
+    @VisibleForTesting
+    static final int ON = 1;
+    @VisibleForTesting
+    static final int OFF = 0;
 
     private SettingObserver mSettingObserver;
 
-    public BadgingNotificationPreferenceController(Context context) {
-        super(context);
+    public BadgingNotificationPreferenceController(Context context, String preferenceKey) {
+        super(context, preferenceKey);
     }
 
     @Override
@@ -78,28 +81,22 @@ public class BadgingNotificationPreferenceController extends AbstractPreferenceC
     }
 
     @Override
-    public String getPreferenceKey() {
-        return KEY_NOTIFICATION_BADGING;
-    }
-
-    @Override
-    public boolean isAvailable() {
+    public int getAvailabilityStatus() {
         return mContext.getResources()
-                .getBoolean(com.android.internal.R.bool.config_notificationBadging);
+                .getBoolean(com.android.internal.R.bool.config_notificationBadging)
+                ? AVAILABLE : DISABLED_UNSUPPORTED;
     }
 
     @Override
-    public void updateState(Preference preference) {
-        final boolean checked = Settings.Secure.getInt(mContext.getContentResolver(),
+    public boolean isChecked() {
+        return Settings.Secure.getInt(mContext.getContentResolver(),
                 NOTIFICATION_BADGING, ON) == ON;
-        ((TwoStatePreference) preference).setChecked(checked);
     }
 
     @Override
-    public boolean onPreferenceChange(Preference preference, Object newValue) {
-        final boolean val = (Boolean) newValue;
+    public boolean setChecked(boolean isChecked) {
         return Settings.Secure.putInt(mContext.getContentResolver(),
-                NOTIFICATION_BADGING, val ? ON : OFF);
+                NOTIFICATION_BADGING, isChecked ? ON : OFF);
     }
 
     class SettingObserver extends ContentObserver {
@@ -134,7 +131,7 @@ public class BadgingNotificationPreferenceController extends AbstractPreferenceC
     @Override
     public ResultPayload getResultPayload() {
         final Intent intent = DatabaseIndexingUtils.buildSearchResultPageIntent(mContext,
-                ConfigureNotificationSettings.class.getName(), KEY_NOTIFICATION_BADGING,
+                ConfigureNotificationSettings.class.getName(), getPreferenceKey(),
                 mContext.getString(R.string.configure_notification_settings));
 
         return new InlineSwitchPayload(Settings.Secure.NOTIFICATION_BADGING,
diff --git a/src/com/android/settings/notification/ConfigureNotificationSettings.java b/src/com/android/settings/notification/ConfigureNotificationSettings.java
index 1ca4a4d6a2c..4bf5b70d7ef 100644
--- a/src/com/android/settings/notification/ConfigureNotificationSettings.java
+++ b/src/com/android/settings/notification/ConfigureNotificationSettings.java
@@ -93,8 +93,6 @@ public class ConfigureNotificationSettings extends DashboardFragment {
     private static List<AbstractPreferenceController> buildPreferenceControllers(Context context,
             Lifecycle lifecycle, Application app, Fragment host) {
         final List<AbstractPreferenceController> controllers = new ArrayList<>();
-        final BadgingNotificationPreferenceController badgeController =
-                new BadgingNotificationPreferenceController(context);
         final PulseNotificationPreferenceController pulseController =
                 new PulseNotificationPreferenceController(context);
         final LockScreenNotificationPreferenceController lockScreenNotificationController =
@@ -108,7 +106,6 @@ public class ConfigureNotificationSettings extends DashboardFragment {
         }
         controllers.add(new RecentNotifyingAppsPreferenceController(
                 context, new NotificationBackend(), app, host));
-        controllers.add(badgeController);
         controllers.add(pulseController);
         controllers.add(lockScreenNotificationController);
         controllers.add(new NotificationRingtonePreferenceController(context) {
diff --git a/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java
index d852aaa5b3e..f9991949e47 100644
--- a/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java
@@ -17,6 +17,12 @@
 package com.android.settings.notification;
 
 import static android.provider.Settings.Secure.NOTIFICATION_BADGING;
+
+import static com.android.settings.core.BasePreferenceController.AVAILABLE;
+import static com.android.settings.core.BasePreferenceController.DISABLED_UNSUPPORTED;
+import static com.android.settings.notification.BadgingNotificationPreferenceController.OFF;
+import static com.android.settings.notification.BadgingNotificationPreferenceController.ON;
+
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -54,10 +60,13 @@ public class BadgingNotificationPreferenceControllerTest {
     private BadgingNotificationPreferenceController mController;
     private Preference mPreference;
 
+    private static final String KEY_NOTIFICATION_BADGING = "notification_badging";
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mController = new BadgingNotificationPreferenceController(mContext);
+        mController = new BadgingNotificationPreferenceController(mContext,
+                KEY_NOTIFICATION_BADGING);
         mPreference = new Preference(RuntimeEnvironment.application);
         mPreference.setKey(mController.getPreferenceKey());
         when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
@@ -88,9 +97,10 @@ public class BadgingNotificationPreferenceControllerTest {
     public void updateState_preferenceSetCheckedWhenSettingIsOn() {
         final TwoStatePreference preference = mock(TwoStatePreference.class);
         final Context context = RuntimeEnvironment.application;
-        Settings.Secure.putInt(context.getContentResolver(), NOTIFICATION_BADGING, 1);
+        Settings.Secure.putInt(context.getContentResolver(), NOTIFICATION_BADGING, ON);
 
-        mController = new BadgingNotificationPreferenceController(context);
+        mController = new BadgingNotificationPreferenceController(context,
+                KEY_NOTIFICATION_BADGING);
         mController.updateState(preference);
 
         verify(preference).setChecked(true);
@@ -100,9 +110,10 @@ public class BadgingNotificationPreferenceControllerTest {
     public void updateState_preferenceSetUncheckedWhenSettingIsOff() {
         final TwoStatePreference preference = mock(TwoStatePreference.class);
         final Context context = RuntimeEnvironment.application;
-        Settings.Secure.putInt(context.getContentResolver(), NOTIFICATION_BADGING, 0);
+        Settings.Secure.putInt(context.getContentResolver(), NOTIFICATION_BADGING, OFF);
 
-        mController = new BadgingNotificationPreferenceController(context);
+        mController = new BadgingNotificationPreferenceController(context,
+                KEY_NOTIFICATION_BADGING);
         mController.updateState(preference);
 
         verify(preference).setChecked(false);
@@ -118,11 +129,11 @@ public class BadgingNotificationPreferenceControllerTest {
     public void testSetValue_updatesCorrectly() {
         final int newValue = 0;
         ContentResolver resolver = mContext.getContentResolver();
-        Settings.Secure.putInt(resolver, Settings.Secure.NOTIFICATION_BADGING, 1);
+        Settings.Secure.putInt(resolver, Settings.Secure.NOTIFICATION_BADGING, ON);
 
         ((InlinePayload) mController.getResultPayload()).setValue(mContext, newValue);
         final int updatedValue =
-            Settings.Secure.getInt(resolver, Settings.Secure.NOTIFICATION_BADGING, 1);
+            Settings.Secure.getInt(resolver, Settings.Secure.NOTIFICATION_BADGING, ON);
 
         assertThat(updatedValue).isEqualTo(newValue);
     }
@@ -138,4 +149,40 @@ public class BadgingNotificationPreferenceControllerTest {
 
         assertThat(newValue).isEqualTo(currentValue);
     }
+
+    @Test
+    public void isChecked_settingIsOff_shouldReturnFalse() {
+        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BADGING, OFF);
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_settingIsOn_shouldReturnTrue() {
+        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BADGING, ON);
+
+        assertThat(mController.isChecked()).isTrue();
+    }
+
+    @Test
+    public void setChecked_setFalse_disablesSetting() {
+        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BADGING, ON);
+
+        mController.setChecked(false);
+        int updatedValue = Settings.Secure.getInt(mContext.getContentResolver(),
+                NOTIFICATION_BADGING, -1);
+
+        assertThat(updatedValue).isEqualTo(OFF);
+    }
+
+    @Test
+    public void setChecked_setTrue_enablesSetting() {
+        Settings.Secure.putInt(mContext.getContentResolver(), NOTIFICATION_BADGING, OFF);
+
+        mController.setChecked(true);
+        int updatedValue = Settings.Secure.getInt(mContext.getContentResolver(),
+                NOTIFICATION_BADGING, -1);
+
+        assertThat(updatedValue).isEqualTo(ON);
+    }
 }