diff --git a/res/xml/ambient_display_settings.xml b/res/xml/ambient_display_settings.xml
index 037421c..306ead5 100644
--- a/res/xml/ambient_display_settings.xml
+++ b/res/xml/ambient_display_settings.xml
@@ -29,8 +29,7 @@
         <SwitchPreference
             android:key="ambient_display_always_on"
             android:title="@string/doze_always_on_title"
-            android:summary="@string/doze_always_on_summary"
-            settings:controller="com.android.settings.display.AmbientDisplayAlwaysOnPreferenceController" />
+            android:summary="@string/doze_always_on_summary" />
 
         <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 e0efaac..9132daa 100644
--- a/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceController.java
+++ b/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceController.java
@@ -19,45 +19,59 @@
 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.TogglePreferenceController;
+import com.android.settings.core.PreferenceControllerMixin;
 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 TogglePreferenceController {
+public class AmbientDisplayAlwaysOnPreferenceController extends
+        AbstractPreferenceController implements PreferenceControllerMixin,
+        Preference.OnPreferenceChangeListener {
 
     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 AmbientDisplayConfiguration mConfig;
-    private OnPreferenceChangedCallback mCallback;
+    private final AmbientDisplayConfiguration mConfig;
+    private final OnPreferenceChangedCallback mCallback;
 
     public interface OnPreferenceChangedCallback {
         void onPreferenceChanged();
     }
 
-    public AmbientDisplayAlwaysOnPreferenceController(Context context, String key) {
-        super(context, key);
+    public AmbientDisplayAlwaysOnPreferenceController(Context context,
+            AmbientDisplayConfiguration config, OnPreferenceChangedCallback callback) {
+        super(context);
+        mConfig = config;
+        mCallback = callback;
     }
 
     @Override
-    public int getAvailabilityStatus() {
-        return isAvailable(mConfig) ? AVAILABLE : DISABLED_UNSUPPORTED;
+    public String getPreferenceKey() {
+        return KEY_ALWAYS_ON;
     }
 
     @Override
-    public boolean isChecked() {
-        return mConfig.alwaysOnEnabled(MY_USER);
+    public void updateState(Preference preference) {
+        ((SwitchPreference) preference).setChecked(isAlwaysOnEnabled(mConfig));
+    }
+
+    public static boolean isAlwaysOnEnabled(AmbientDisplayConfiguration config) {
+        return config.alwaysOnEnabled(MY_USER);
     }
 
     @Override
-    public boolean setChecked(boolean isChecked) {
-        int enabled = isChecked ? ON : OFF;
+    public boolean onPreferenceChange(Preference preference, Object newValue) {
+        int enabled = (boolean) newValue ? ON : OFF;
         Settings.Secure.putInt(
                 mContext.getContentResolver(), Settings.Secure.DOZE_ALWAYS_ON, enabled);
         if (mCallback != null) {
@@ -66,16 +80,9 @@
         return true;
     }
 
-    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);
+    @Override
+    public boolean isAvailable() {
+        return alwaysOnAvailableForUser(mConfig);
     }
 
     public static boolean isAvailable(AmbientDisplayConfiguration config) {
@@ -89,11 +96,16 @@
     @Override
     public ResultPayload getResultPayload() {
         final Intent intent = DatabaseIndexingUtils.buildSearchResultPageIntent(mContext,
-                AmbientDisplaySettings.class.getName(), getPreferenceKey(),
+                AmbientDisplaySettings.class.getName(), KEY_ALWAYS_ON,
                 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 a88d2ed..02db6f1 100644
--- a/src/com/android/settings/display/AmbientDisplaySettings.java
+++ b/src/com/android/settings/display/AmbientDisplaySettings.java
@@ -40,23 +40,23 @@
  */
 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) {
+            MetricsFeatureProvider metricsFeatureProvider,
+            AmbientDisplayAlwaysOnPreferenceController.OnPreferenceChangedCallback aodCallback) {
         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,14 +64,6 @@
         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() {
@@ -86,7 +78,8 @@
     @Override
     protected List<AbstractPreferenceController> createPreferenceControllers(Context context) {
         return buildPreferenceControllers(context, getLifecycle(),
-                getConfig(context), mMetricsFeatureProvider);
+                new AmbientDisplayConfiguration(context), mMetricsFeatureProvider,
+                this::updatePreferenceStates);
     }
 
     @Override
@@ -111,14 +104,7 @@
                 public List<AbstractPreferenceController> createPreferenceControllers(
                         Context context) {
                     return buildPreferenceControllers(context, null,
-                            new AmbientDisplayConfiguration(context), null);
+                            new AmbientDisplayConfiguration(context), null, 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 e1c4b70..d06ea2a 100644
--- a/tests/robotests/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/AmbientDisplayAlwaysOnPreferenceControllerTest.java
@@ -17,14 +17,17 @@
 package com.android.settings.display;
 
 import static com.google.common.truth.Truth.assertThat;
-
-import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Mockito.doReturn;
 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;
@@ -46,6 +49,8 @@
 
     @Mock
     private AmbientDisplayConfiguration mConfig;
+    @Mock
+    private SwitchPreference mSwitchPreference;
 
     private Context mContext;
 
@@ -59,90 +64,95 @@
         MockitoAnnotations.initMocks(this);
         mContext = RuntimeEnvironment.application;
         mContentResolver = mContext.getContentResolver();
-        mController = new AmbientDisplayAlwaysOnPreferenceController(mContext, "key");
-        mController.setConfig(mConfig);
-        mController.setCallback(() -> mCallbackInvoked = true);
+        mController = new AmbientDisplayAlwaysOnPreferenceController(mContext, mConfig,
+                () -> {
+                    mCallbackInvoked = true;
+                });
     }
 
     @Test
-    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() {
+    public void updateState_enabled() {
         when(mConfig.alwaysOnEnabled(anyInt())).thenReturn(true);
 
-        assertThat(mController.isChecked()).isTrue();
+        mController.updateState(mSwitchPreference);
+
+        verify(mSwitchPreference).setChecked(true);
     }
 
     @Test
-    public void isChecked_disabled() {
+    public void updateState_disabled() {
         when(mConfig.alwaysOnEnabled(anyInt())).thenReturn(false);
 
-        assertThat(mController.isChecked()).isFalse();
-    }
+        mController.updateState(mSwitchPreference);
 
-    @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);
+        verify(mSwitchPreference).setChecked(false);
     }
 
     @Test
     public void onPreferenceChange_callback() {
         assertThat(mCallbackInvoked).isFalse();
-        mController.setChecked(true);
+        mController.onPreferenceChange(mSwitchPreference, true);
         assertThat(mCallbackInvoked).isTrue();
     }
 
     @Test
-    public void testPreferenceController_ProperResultPayloadType() {
-        when(mConfig.alwaysOnAvailableForUser(anyInt())).thenReturn(false);
+    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() {
+        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
deleted file mode 100644
index bab5d41..0000000
--- a/tests/robotests/src/com/android/settings/display/AmbientDisplaySettingsTest.java
+++ /dev/null
@@ -1,69 +0,0 @@
-/*
- * 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.
- */
-
-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_shouldSetConfigAndCallback() {
-        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