diff --git a/src/com/android/settings/dashboard/DashboardFragment.java b/src/com/android/settings/dashboard/DashboardFragment.java
index a730939..ef49e30 100644
--- a/src/com/android/settings/dashboard/DashboardFragment.java
+++ b/src/com/android/settings/dashboard/DashboardFragment.java
@@ -120,7 +120,8 @@
     @Override
     public void notifySummaryChanged(Tile tile) {
         final String key = mDashboardFeatureProvider.getDashboardKeyForTile(tile);
-        final Preference pref = findPreference(key);
+        final Preference pref = mProgressiveDisclosureMixin.findPreference(
+                getPreferenceScreen(), key);
         if (pref == null) {
             Log.d(getLogTag(),
                     String.format("Can't find pref by key %s, skipping update summary %s/%s",
@@ -164,18 +165,6 @@
         }
     }
 
-    @Override
-    public Preference findPreference(CharSequence key) {
-        Preference preference = super.findPreference(key);
-        if (preference == null && mProgressiveDisclosureMixin != null) {
-            preference = mProgressiveDisclosureMixin.findPreference(key);
-        }
-        if (preference == null) {
-            Log.d(TAG, "Cannot find preference with key " + key);
-        }
-        return preference;
-    }
-
     protected <T extends PreferenceController> T getPreferenceController(Class<T> clazz) {
         PreferenceController controller = mPreferenceControllers.get(clazz);
         return (T) controller;
@@ -221,10 +210,78 @@
         }
     }
 
+    @Override
+    public View onCreateView(LayoutInflater inflater, ViewGroup container,
+            Bundle savedInstanceState) {
+        final View view = super.onCreateView(inflater, container, savedInstanceState);
+        if (mDashboardFeatureProvider.isEnabled()) {
+            getListView().addItemDecoration(mDividerDecoration);
+        }
+        return view;
+    }
+
     /**
-     * Displays dashboard tiles as preference.
+     * Update state of each preference managed by PreferenceController.
      */
-    private final void displayDashboardTiles(final String TAG, PreferenceScreen screen) {
+    private void updatePreferenceStates() {
+        Collection<PreferenceController> controllers = mPreferenceControllers.values();
+        final PreferenceScreen screen = getPreferenceScreen();
+        for (PreferenceController controller : controllers) {
+            final String key = controller.getPreferenceKey();
+
+            final Preference preference = mProgressiveDisclosureMixin.findPreference(screen, key);
+            if (preference == null) {
+                Log.d(TAG, "Cannot find preference with key " + key);
+                continue;
+            }
+            controller.updateState(preference);
+        }
+    }
+
+    @Override
+    public void setDivider(Drawable divider) {
+        if (mDashboardFeatureProvider.isEnabled()) {
+            // Intercept divider and set it transparent so system divider decoration is disabled.
+            // We will use our decoration to draw divider more intelligently.
+            mDividerDecoration.setDivider(divider);
+            super.setDivider(new ColorDrawable(Color.TRANSPARENT));
+        } else {
+            super.setDivider(divider);
+        }
+    }
+
+    /**
+     * Refresh all preference items, including both static prefs from xml, and dynamic items from
+     * DashboardCategory.
+     */
+    private void refreshAllPreferences(final String TAG) {
+        // First remove old preferences.
+        if (getPreferenceScreen() != null) {
+            // Intentionally do not cache PreferenceScreen because it will be recreated later.
+            getPreferenceScreen().removeAll();
+        }
+
+        // Add resource based tiles.
+        displayResourceTiles();
+
+        refreshDashboardTiles(TAG);
+
+        if (!mProgressiveDisclosureMixin.isCollapsed()
+                && mProgressiveDisclosureMixin.shouldCollapse(getPreferenceScreen())) {
+            mProgressiveDisclosureMixin.collapse(getPreferenceScreen());
+        }
+    }
+
+    /**
+     * Refresh preference items backed by DashboardCategory.
+     */
+    private void refreshDashboardTiles(final String TAG) {
+        final PreferenceScreen screen = getPreferenceScreen();
+        for (String key : mDashboardTilePrefKeys) {
+            // Remove tiles from screen
+            mProgressiveDisclosureMixin.removePreference(screen, key);
+        }
+        mDashboardTilePrefKeys.clear();
         final Context context = getContext();
         final DashboardCategory category =
                 mDashboardFeatureProvider.getTilesForCategory(getCategoryKey());
@@ -275,98 +332,7 @@
             // (larger value has higher priority). However pref order defines smaller value has
             // higher priority.
             pref.setOrder(-tile.priority);
-
-            // Either add to screen, or to collapsed list.
-            if (mProgressiveDisclosureMixin.isCollapsed()) {
-                // Already collapsed, add to collapsed list.
-                mProgressiveDisclosureMixin.addToCollapsedList(pref);
-            } else if (mProgressiveDisclosureMixin.shouldCollapse(screen)) {
-                // About to have too many tiles on scree, collapse and add pref to collapsed list.
-                mProgressiveDisclosureMixin.collapse(screen);
-                mProgressiveDisclosureMixin.addToCollapsedList(pref);
-            } else {
-                // No need to collapse, add to screen directly.
-                screen.addPreference(pref);
-            }
+            mProgressiveDisclosureMixin.addPreference(screen, pref);
         }
     }
-
-    @Override
-    public View onCreateView(LayoutInflater inflater, ViewGroup container,
-            Bundle savedInstanceState) {
-        final View view = super.onCreateView(inflater, container, savedInstanceState);
-        if (mDashboardFeatureProvider.isEnabled()) {
-            getListView().addItemDecoration(mDividerDecoration);
-        }
-        return view;
-    }
-
-    /**
-     * Update state of each preference managed by PreferenceController.
-     */
-    private void updatePreferenceStates() {
-        Collection<PreferenceController> controllers = mPreferenceControllers.values();
-        for (PreferenceController controller : controllers) {
-            final String key = controller.getPreferenceKey();
-
-            final Preference preference = findPreference(key);
-            if (preference == null) {
-                Log.d(TAG, "Cannot find preference with key " + key);
-                continue;
-            }
-            controller.updateState(preference);
-        }
-    }
-
-    @Override
-    public void setDivider(Drawable divider) {
-        if (mDashboardFeatureProvider.isEnabled()) {
-            // Intercept divider and set it transparent so system divider decoration is disabled.
-            // We will use our decoration to draw divider more intelligently.
-            mDividerDecoration.setDivider(divider);
-            super.setDivider(new ColorDrawable(Color.TRANSPARENT));
-        } else {
-            super.setDivider(divider);
-        }
-    }
-
-    /**
-     * Refresh all preference items, including both static prefs from xml, and dynamic items from
-     * DashboardCategory.
-     */
-    private void refreshAllPreferences(final String TAG) {
-        // First remove old preferences.
-        if (getPreferenceScreen() != null) {
-            // Intentionally do not cache PreferenceScreen because it will be recreated later.
-            getPreferenceScreen().removeAll();
-        }
-
-        // Add resource based tiles.
-        displayResourceTiles();
-
-        refreshDashboardTiles(TAG);
-
-        if (!mProgressiveDisclosureMixin.isCollapsed()
-                && mProgressiveDisclosureMixin.shouldCollapse(getPreferenceScreen())) {
-            mProgressiveDisclosureMixin.collapse(getPreferenceScreen());
-        }
-    }
-
-    /**
-     * Refresh preference items backed by DashboardCategory.
-     */
-    private void refreshDashboardTiles(final String TAG) {
-        final PreferenceScreen screen = getPreferenceScreen();
-        for (String key : mDashboardTilePrefKeys) {
-            // Remove tiles from screen
-            final Preference pref = screen.findPreference(key);
-            if (pref != null) {
-                screen.removePreference(pref);
-            }
-            // Also remove tile from collapsed set
-            mProgressiveDisclosureMixin.removePreference(screen, key);
-        }
-        mDashboardTilePrefKeys.clear();
-        displayDashboardTiles(TAG, getPreferenceScreen());
-    }
 }
diff --git a/src/com/android/settings/dashboard/ProgressiveDisclosureMixin.java b/src/com/android/settings/dashboard/ProgressiveDisclosureMixin.java
index 28113ad..84b6b7f 100644
--- a/src/com/android/settings/dashboard/ProgressiveDisclosureMixin.java
+++ b/src/com/android/settings/dashboard/ProgressiveDisclosureMixin.java
@@ -18,6 +18,7 @@
 
 import android.content.Context;
 import android.os.Bundle;
+import android.support.annotation.VisibleForTesting;
 import android.support.v14.preference.PreferenceFragment;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
@@ -77,6 +78,7 @@
                 for (Preference pref : collapsedPrefs) {
                     screen.addPreference(pref);
                 }
+                collapsedPrefs.clear();
                 mUserExpanded = true;
             }
         }
@@ -126,19 +128,36 @@
     }
 
     /**
-     * Add preference to collapsed list.
+     * Adds preference to screen. If there are too many preference on screen, adds it to
+     * collapsed list instead.
      */
-    public void addToCollapsedList(Preference preference) {
-        collapsedPrefs.add(preference);
+    public void addPreference(PreferenceScreen screen, Preference pref) {
+        // Either add to screen, or to collapsed list.
+        if (isCollapsed()) {
+            // Already collapsed, add to collapsed list.
+            addToCollapsedList(pref);
+        } else if (shouldCollapse(screen)) {
+            // About to have too many tiles on scree, collapse and add pref to collapsed list.
+            collapse(screen);
+            addToCollapsedList(pref);
+        } else {
+            // No need to collapse, add to screen directly.
+            screen.addPreference(pref);
+        }
     }
 
     /**
-     * Remove preference from collapsed list. If the preference is not in list, do nothing.
+     * Removes preference. If the preference is on screen, remove it from screen. If the
+     * preference is in collapsed list, remove it from list.
      */
     public void removePreference(PreferenceScreen screen, String key) {
-        if (!isCollapsed()) {
+        // Try removing from screen.
+        final Preference preference = screen.findPreference(key);
+        if (preference != null) {
+            screen.removePreference(preference);
             return;
         }
+        // Didn't find on screen, try removing from collapsed list.
         for (int i = 0; i < collapsedPrefs.size(); i++) {
             final Preference pref = collapsedPrefs.get(i);
             if (TextUtils.equals(key, pref.getKey())) {
@@ -153,16 +172,29 @@
     }
 
     /**
-     * Find whether a preference is in collapsed list.
+     * Finds preference by key, either from screen or from collapsed list.
      */
-    public Preference findPreference(CharSequence key) {
+    public Preference findPreference(PreferenceScreen screen, CharSequence key) {
+        Preference preference = screen.findPreference(key);
+        if (preference != null) {
+            return preference;
+        }
         for (int i = 0; i < collapsedPrefs.size(); i++) {
             final Preference pref = collapsedPrefs.get(i);
             if (TextUtils.equals(key, pref.getKey())) {
                 return pref;
             }
         }
+        Log.d(TAG, "Cannot find preference with key " + key);
         return null;
     }
 
+    /**
+     * Add preference to collapsed list.
+     */
+    @VisibleForTesting
+    void addToCollapsedList(Preference preference) {
+        collapsedPrefs.add(preference);
+    }
+
 }
diff --git a/tests/robotests/src/com/android/settings/dashboard/ProgressiveDisclosureTest.java b/tests/robotests/src/com/android/settings/dashboard/ProgressiveDisclosureTest.java
index 6964d74..986ccd9 100644
--- a/tests/robotests/src/com/android/settings/dashboard/ProgressiveDisclosureTest.java
+++ b/tests/robotests/src/com/android/settings/dashboard/ProgressiveDisclosureTest.java
@@ -38,6 +38,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -53,6 +54,7 @@
     private FakeFeatureFactory mFakeFeatureFactory;
     @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     private PreferenceFragment mPreferenceFragment;
+    private PreferenceScreen mScreen;
     private Context mAppContext;
     private Preference mPreference;
 
@@ -62,6 +64,7 @@
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         FakeFeatureFactory.setupForTest(mContext);
+        mScreen = mPreferenceFragment.getPreferenceScreen();
         mAppContext = ShadowApplication.getInstance().getApplicationContext();
         mFakeFeatureFactory = (FakeFeatureFactory) FeatureFactory.getFactory(mContext);
         mMixin = new ProgressiveDisclosureMixin(mAppContext, mPreferenceFragment);
@@ -72,57 +75,81 @@
 
     @Test
     public void shouldNotCollapse_lessPreferenceThanLimit() {
-        when(mPreferenceFragment.getPreferenceScreen().getPreferenceCount()).thenReturn(5);
+        when(mScreen.getPreferenceCount()).thenReturn(5);
 
         mMixin.setTileLimit(10);
 
-        assertThat(mMixin.shouldCollapse(mPreferenceFragment.getPreferenceScreen())).isFalse();
+        assertThat(mMixin.shouldCollapse(mScreen)).isFalse();
     }
 
     @Test
     public void shouldCollapse_morePreferenceThanLimit() {
         when(mFakeFeatureFactory.dashboardFeatureProvider.isEnabled()).thenReturn(true);
-        when(mPreferenceFragment.getPreferenceScreen().getPreferenceCount()).thenReturn(5);
+        when(mScreen.getPreferenceCount()).thenReturn(5);
 
-        assertThat(mMixin.shouldCollapse(mPreferenceFragment.getPreferenceScreen())).isTrue();
+        assertThat(mMixin.shouldCollapse(mScreen)).isTrue();
     }
 
     @Test
     public void findPreference_prefInCollapsedList_shouldFindIt() {
+        when(mScreen.findPreference(anyString())).thenReturn(null);
         mMixin.addToCollapsedList(mPreference);
 
-        Preference pref = mMixin.findPreference(mPreference.getKey());
+        Preference pref = mMixin.findPreference(mScreen, mPreference.getKey());
 
         assertThat(pref).isNotNull();
         assertThat(pref).isSameAs(mPreference);
     }
 
     @Test
-    public void findPreference_prefNotInCollapsedList_shouldNotFindIt() {
-        Preference pref = mMixin.findPreference(mPreference.getKey());
+    public void findPreference_prefOnScreen_shouldFindIt() {
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
+
+        Preference pref = mMixin.findPreference(mScreen, mPreference.getKey());
+
+        assertThat(pref).isNotNull();
+        assertThat(pref).isSameAs(mPreference);
+    }
+
+    @Test
+    public void findPreference_prefNotInCollapsedListOrScreen_shouldNotFindIt() {
+        when(mScreen.findPreference(anyString())).thenReturn(null);
+        Preference pref = mMixin.findPreference(mScreen, mPreference.getKey());
 
         assertThat(pref).isNull();
     }
 
     @Test
     public void findPreference_prefRemovedFromCollapsedList_shouldNotFindIt() {
+        when(mScreen.findPreference(anyString())).thenReturn(null);
         mMixin.addToCollapsedList(mPreference);
         mMixin.removePreference(mPreferenceFragment.getPreferenceScreen(), mPreference.getKey());
-        Preference pref = mMixin.findPreference(mPreference.getKey());
+
+        Preference pref = mMixin.findPreference(mScreen, mPreference.getKey());
 
         assertThat(pref).isNull();
     }
 
     @Test
+    public void removePreference_shouldRemoveOnScreenPreference() {
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
+
+        mMixin.removePreference(mScreen, mPreference.getKey());
+
+        verify(mScreen).removePreference(mPreference);
+    }
+
+    @Test
     public void removeLastPreference_shouldRemoveExpandButtonToo() {
+        when(mScreen.findPreference(anyString())).thenReturn(null);
         mMixin.addToCollapsedList(mPreference);
         // Collapsed
         assertThat(mMixin.isCollapsed()).isTrue();
 
-        mMixin.removePreference(mPreferenceFragment.getPreferenceScreen(), mPreference.getKey());
+        mMixin.removePreference(mScreen, mPreference.getKey());
 
         // Removing expand button
-        verify(mPreferenceFragment.getPreferenceScreen()).removePreference(any(Preference.class));
+        verify(mScreen).removePreference(any(Preference.class));
         // No longer collapsed
         assertThat(mMixin.isCollapsed()).isFalse();
     }