diff --git a/res/xml/ambient_display_settings.xml b/res/xml/ambient_display_settings.xml
index 306ead5d50f..037421cd213 100644
--- a/res/xml/ambient_display_settings.xml
+++ b/res/xml/ambient_display_settings.xml
@@ -29,7 +29,8 @@
         <SwitchPreference
             android:key="ambient_display_always_on"
             android:title="@string/doze_always_on_title"
-            android:summary="@string/doze_always_on_summary" />
+            android:summary="@string/doze_always_on_summary"
+            settings:controller="com.android.settings.display.AmbientDisplayAlwaysOnPreferenceController" />
 
         <Preference
             android:key="ambient_display_double_tap"
diff --git a/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceController.java b/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceController.java
index 9132daa691c..e0efaacb87c 100644
--- a/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceController.java
+++ b/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceController.java
@@ -19,59 +19,45 @@ import android.content.Context;
 import android.content.Intent;
 import android.os.UserHandle;
 import android.provider.Settings;
-import android.support.v14.preference.SwitchPreference;
-import android.support.v7.preference.Preference;
-import android.support.annotation.VisibleForTesting;
 
 import com.android.internal.hardware.AmbientDisplayConfiguration;
 import com.android.settings.R;
-import com.android.settings.core.PreferenceControllerMixin;
+import com.android.settings.core.TogglePreferenceController;
 import com.android.settings.search.DatabaseIndexingUtils;
 import com.android.settings.search.InlineSwitchPayload;
 import com.android.settings.search.ResultPayload;
-import com.android.settingslib.core.AbstractPreferenceController;
 
-public class AmbientDisplayAlwaysOnPreferenceController extends
-        AbstractPreferenceController implements PreferenceControllerMixin,
-        Preference.OnPreferenceChangeListener {
+public class AmbientDisplayAlwaysOnPreferenceController extends TogglePreferenceController {
 
     private final int ON = 1;
     private final int OFF = 0;
 
-    public static final String KEY_ALWAYS_ON = "ambient_display_always_on";
     private static final int MY_USER = UserHandle.myUserId();
 
-    private final AmbientDisplayConfiguration mConfig;
-    private final OnPreferenceChangedCallback mCallback;
+    private AmbientDisplayConfiguration mConfig;
+    private OnPreferenceChangedCallback mCallback;
 
     public interface OnPreferenceChangedCallback {
         void onPreferenceChanged();
     }
 
-    public AmbientDisplayAlwaysOnPreferenceController(Context context,
-            AmbientDisplayConfiguration config, OnPreferenceChangedCallback callback) {
-        super(context);
-        mConfig = config;
-        mCallback = callback;
+    public AmbientDisplayAlwaysOnPreferenceController(Context context, String key) {
+        super(context, key);
     }
 
     @Override
-    public String getPreferenceKey() {
-        return KEY_ALWAYS_ON;
+    public int getAvailabilityStatus() {
+        return isAvailable(mConfig) ? AVAILABLE : DISABLED_UNSUPPORTED;
     }
 
     @Override
-    public void updateState(Preference preference) {
-        ((SwitchPreference) preference).setChecked(isAlwaysOnEnabled(mConfig));
-    }
-
-    public static boolean isAlwaysOnEnabled(AmbientDisplayConfiguration config) {
-        return config.alwaysOnEnabled(MY_USER);
+    public boolean isChecked() {
+        return mConfig.alwaysOnEnabled(MY_USER);
     }
 
     @Override
-    public boolean onPreferenceChange(Preference preference, Object newValue) {
-        int enabled = (boolean) newValue ? ON : OFF;
+    public boolean setChecked(boolean isChecked) {
+        int enabled = isChecked ? ON : OFF;
         Settings.Secure.putInt(
                 mContext.getContentResolver(), Settings.Secure.DOZE_ALWAYS_ON, enabled);
         if (mCallback != null) {
@@ -80,9 +66,16 @@ public class AmbientDisplayAlwaysOnPreferenceController extends
         return true;
     }
 
-    @Override
-    public boolean isAvailable() {
-        return alwaysOnAvailableForUser(mConfig);
+    public void setConfig(AmbientDisplayConfiguration config) {
+        mConfig = config;
+    }
+
+    public void setCallback(OnPreferenceChangedCallback callback) {
+        mCallback = callback;
+    }
+
+    public static boolean isAlwaysOnEnabled(AmbientDisplayConfiguration config) {
+        return config.alwaysOnEnabled(MY_USER);
     }
 
     public static boolean isAvailable(AmbientDisplayConfiguration config) {
@@ -96,16 +89,11 @@ public class AmbientDisplayAlwaysOnPreferenceController extends
     @Override
     public ResultPayload getResultPayload() {
         final Intent intent = DatabaseIndexingUtils.buildSearchResultPageIntent(mContext,
-                AmbientDisplaySettings.class.getName(), KEY_ALWAYS_ON,
+                AmbientDisplaySettings.class.getName(), getPreferenceKey(),
                 mContext.getString(R.string.ambient_display_screen_title));
 
         return new InlineSwitchPayload(Settings.Secure.DOZE_ALWAYS_ON,
                 ResultPayload.SettingsSource.SECURE, ON /* onValue */, intent, isAvailable(),
                 ON /* defaultValue */);
     }
-
-    @VisibleForTesting
-    boolean alwaysOnAvailableForUser(AmbientDisplayConfiguration config) {
-        return isAvailable(config);
-    }
 }
diff --git a/src/com/android/settings/display/AmbientDisplaySettings.java b/src/com/android/settings/display/AmbientDisplaySettings.java
index 02db6f1a4c0..db3f38ae165 100644
--- a/src/com/android/settings/display/AmbientDisplaySettings.java
+++ b/src/com/android/settings/display/AmbientDisplaySettings.java
@@ -40,23 +40,23 @@ import java.util.List;
  */
 public class AmbientDisplaySettings extends DashboardFragment {
 
+    public static final String KEY_AMBIENT_DISPLAY_ALWAYS_ON = "ambient_display_always_on";
+
     private static final String TAG = "AmbientDisplaySettings";
     private static final int MY_USER_ID = UserHandle.myUserId();
 
-    private static final String KEY_AMBIENT_DISPLAY_ALWAYS_ON = "ambient_display_always_on";
     private static final String KEY_AMBIENT_DISPLAY_DOUBLE_TAP = "ambient_display_double_tap";
     private static final String KEY_AMBIENT_DISPLAY_PICK_UP = "ambient_display_pick_up";
     private static final String KEY_AMBIENT_DISPLAY_NOTIFICATION = "ambient_display_notification";
 
+    private AmbientDisplayConfiguration mConfig;
+
     private static List<AbstractPreferenceController> buildPreferenceControllers(Context context,
             Lifecycle lifecycle, AmbientDisplayConfiguration config,
-            MetricsFeatureProvider metricsFeatureProvider,
-            AmbientDisplayAlwaysOnPreferenceController.OnPreferenceChangedCallback aodCallback) {
+            MetricsFeatureProvider metricsFeatureProvider) {
         final List<AbstractPreferenceController> controllers = new ArrayList<>();
         controllers.add(new AmbientDisplayNotificationsPreferenceController(context, config,
                 metricsFeatureProvider));
-        controllers.add(new AmbientDisplayAlwaysOnPreferenceController(context, config,
-                aodCallback));
         controllers.add(new DoubleTapScreenPreferenceController(context, lifecycle, config,
                 MY_USER_ID, KEY_AMBIENT_DISPLAY_DOUBLE_TAP));
         controllers.add(new PickupGesturePreferenceController(context, lifecycle, config,
@@ -64,6 +64,14 @@ public class AmbientDisplaySettings extends DashboardFragment {
         return controllers;
     }
 
+    @Override
+    public void onAttach(Context context) {
+        super.onAttach(context);
+        final AmbientDisplayAlwaysOnPreferenceController controller = use(
+                AmbientDisplayAlwaysOnPreferenceController.class);
+        controller.setConfig(getConfig(context));
+        controller.setCallback(this::updatePreferenceStates);
+    }
 
     @Override
     protected String getLogTag() {
@@ -78,8 +86,7 @@ public class AmbientDisplaySettings extends DashboardFragment {
     @Override
     protected List<AbstractPreferenceController> createPreferenceControllers(Context context) {
         return buildPreferenceControllers(context, getLifecycle(),
-                new AmbientDisplayConfiguration(context), mMetricsFeatureProvider,
-                this::updatePreferenceStates);
+                getConfig(context), mMetricsFeatureProvider);
     }
 
     @Override
@@ -104,7 +111,14 @@ public class AmbientDisplaySettings extends DashboardFragment {
                 public List<AbstractPreferenceController> createPreferenceControllers(
                         Context context) {
                     return buildPreferenceControllers(context, null,
-                            new AmbientDisplayConfiguration(context), null, null);
+                            new AmbientDisplayConfiguration(context), null);
                 }
             };
+
+    private AmbientDisplayConfiguration getConfig(Context context) {
+        if (mConfig == null) {
+            mConfig = new AmbientDisplayConfiguration(context);
+        }
+        return mConfig;
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceControllerTest.java
index d06ea2ac50c..e1c4b70477a 100644
--- a/tests/robotests/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceControllerTest.java
@@ -17,17 +17,14 @@
 package com.android.settings.display;
 
 import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.doReturn;
+
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.ContentResolver;
 import android.content.Context;
 import android.provider.Settings;
-import android.support.v14.preference.SwitchPreference;
 
 import com.android.internal.hardware.AmbientDisplayConfiguration;
 import com.android.settings.search.InlinePayload;
@@ -49,8 +46,6 @@ public class AmbientDisplayAlwaysOnPreferenceControllerTest {
 
     @Mock
     private AmbientDisplayConfiguration mConfig;
-    @Mock
-    private SwitchPreference mSwitchPreference;
 
     private Context mContext;
 
@@ -64,95 +59,90 @@ public class AmbientDisplayAlwaysOnPreferenceControllerTest {
         MockitoAnnotations.initMocks(this);
         mContext = RuntimeEnvironment.application;
         mContentResolver = mContext.getContentResolver();
-        mController = new AmbientDisplayAlwaysOnPreferenceController(mContext, mConfig,
-                () -> {
-                    mCallbackInvoked = true;
-                });
+        mController = new AmbientDisplayAlwaysOnPreferenceController(mContext, "key");
+        mController.setConfig(mConfig);
+        mController.setCallback(() -> mCallbackInvoked = true);
     }
 
     @Test
-    public void updateState_enabled() {
+    public void getAvailabilityStatus_available() {
+        when(mConfig.alwaysOnAvailableForUser(anyInt())).thenReturn(true);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(
+                AmbientDisplayAlwaysOnPreferenceController.AVAILABLE);
+    }
+
+    @Test
+    public void getAvailabilityStatus_disabled_unsupported() {
+        when(mConfig.alwaysOnAvailableForUser(anyInt())).thenReturn(false);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(
+                AmbientDisplayAlwaysOnPreferenceController.DISABLED_UNSUPPORTED);
+    }
+
+    @Test
+    public void isChecked_enabled() {
         when(mConfig.alwaysOnEnabled(anyInt())).thenReturn(true);
 
-        mController.updateState(mSwitchPreference);
-
-        verify(mSwitchPreference).setChecked(true);
+        assertThat(mController.isChecked()).isTrue();
     }
 
     @Test
-    public void updateState_disabled() {
+    public void isChecked_disabled() {
         when(mConfig.alwaysOnEnabled(anyInt())).thenReturn(false);
 
-        mController.updateState(mSwitchPreference);
+        assertThat(mController.isChecked()).isFalse();
+    }
 
-        verify(mSwitchPreference).setChecked(false);
+    @Test
+    public void setChecked_enabled() {
+        mController.setChecked(true);
+
+        assertThat(Settings.Secure.getInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, -1))
+                .isEqualTo(1);
+    }
+
+    @Test
+    public void setChecked_disabled() {
+        mController.setChecked(false);
+
+        assertThat(Settings.Secure.getInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, -1))
+                .isEqualTo(0);
     }
 
     @Test
     public void onPreferenceChange_callback() {
         assertThat(mCallbackInvoked).isFalse();
-        mController.onPreferenceChange(mSwitchPreference, true);
+        mController.setChecked(true);
         assertThat(mCallbackInvoked).isTrue();
     }
 
-    @Test
-    public void onPreferenceChange_enable() {
-        mController.onPreferenceChange(mSwitchPreference, true);
-
-        assertThat(Settings.Secure.getInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, -1))
-            .isEqualTo(1);
-    }
-
-    @Test
-    public void onPreferenceChange_disable() {
-        mController.onPreferenceChange(mSwitchPreference, false);
-
-        assertThat(Settings.Secure.getInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, -1))
-            .isEqualTo(0);
-    }
-
-    @Test
-    public void isAvailable_available() {
-        mController = spy(mController);
-        doReturn(true).when(mController).alwaysOnAvailableForUser(any());
-
-        assertThat(mController.isAvailable()).isTrue();
-    }
-
-    @Test
-    public void isAvailable_unavailable() {
-        mController = spy(mController);
-        doReturn(false).when(mController).alwaysOnAvailableForUser(any());
-
-        assertThat(mController.isAvailable()).isFalse();
-    }
-
     @Test
     public void testPreferenceController_ProperResultPayloadType() {
+        when(mConfig.alwaysOnAvailableForUser(anyInt())).thenReturn(false);
         mController = spy(mController);
-        doReturn(false).when(mController).alwaysOnAvailableForUser(any());
 
         assertThat(mController.getResultPayload()).isInstanceOf(InlineSwitchPayload.class);
     }
 
     @Test
     public void testSetValue_updatesCorrectly() {
+        when(mConfig.alwaysOnAvailableForUser(anyInt())).thenReturn(false);
         mController = spy(mController);
-        doReturn(false).when(mController).alwaysOnAvailableForUser(any());
         final int newValue = 1;
         Settings.Secure.putInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, 0 /* value */);
 
         ((InlinePayload) mController.getResultPayload()).setValue(mContext, newValue);
         final int updatedValue = Settings.Secure.
-            getInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, 1 /* default */);
+                getInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, 1 /* default */);
 
         assertThat(updatedValue).isEqualTo(newValue);
     }
 
     @Test
     public void testGetValue_correctValueReturned() {
+        when(mConfig.alwaysOnAvailableForUser(anyInt())).thenReturn(false);
         mController = spy(mController);
-        doReturn(false).when(mController).alwaysOnAvailableForUser(any());
         final int currentValue = 1;
         Settings.Secure.putInt(mContentResolver, Settings.Secure.DOZE_ALWAYS_ON, currentValue);
 
diff --git a/tests/robotests/src/com/android/settings/display/AmbientDisplaySettingsTest.java b/tests/robotests/src/com/android/settings/display/AmbientDisplaySettingsTest.java
new file mode 100644
index 00000000000..2f6125294c5
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/display/AmbientDisplaySettingsTest.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.display;
+
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+import com.android.settingslib.core.AbstractPreferenceController;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(SettingsRobolectricTestRunner.class)
+public class AmbientDisplaySettingsTest {
+
+    private TestFragment mTestFragment;
+
+    private Context mContext;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mTestFragment = spy(new TestFragment());
+    }
+
+    @Test
+    public void onAttach_shouldInvokeSetters() {
+        final AmbientDisplayAlwaysOnPreferenceController controller = mock(
+                AmbientDisplayAlwaysOnPreferenceController.class);
+        doReturn(controller).when(mTestFragment).use(
+                AmbientDisplayAlwaysOnPreferenceController.class);
+
+        mTestFragment.onAttach(mContext);
+
+        verify(controller).setConfig(any());
+        verify(controller).setCallback(any());
+    }
+
+    public static class TestFragment extends AmbientDisplaySettings {
+        @Override
+        protected <T extends AbstractPreferenceController> T use(Class<T> clazz) {
+            return super.use(clazz);
+        }
+    }
+}
\ No newline at end of file