Merge "Use setVisible instead of removePreference everywhere."
diff --git a/src/com/android/settings/deviceinfo/SystemUpdatePreferenceController.java b/src/com/android/settings/deviceinfo/SystemUpdatePreferenceController.java
index d2ad6d8..d8a64a8 100644
--- a/src/com/android/settings/deviceinfo/SystemUpdatePreferenceController.java
+++ b/src/com/android/settings/deviceinfo/SystemUpdatePreferenceController.java
@@ -15,6 +15,8 @@
  */
 package com.android.settings.deviceinfo;
 
+import static android.content.Context.CARRIER_CONFIG_SERVICE;
+
 import android.content.Context;
 import android.content.Intent;
 import android.os.Build;
@@ -31,8 +33,6 @@
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settingslib.core.AbstractPreferenceController;
 
-import static android.content.Context.CARRIER_CONFIG_SERVICE;
-
 public class SystemUpdatePreferenceController extends AbstractPreferenceController implements
         PreferenceControllerMixin {
 
@@ -59,12 +59,11 @@
 
     @Override
     public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
         if (isAvailable()) {
             Utils.updatePreferenceToSpecificActivityOrRemove(mContext, screen,
                     KEY_SYSTEM_UPDATE_SETTINGS,
                     Utils.UPDATE_PREFERENCE_FLAG_SET_TITLE_TO_MATCHING_ACTIVITY);
-        } else {
-            removePreference(screen, KEY_SYSTEM_UPDATE_SETTINGS);
         }
     }
 
diff --git a/src/com/android/settings/display/NightModePreferenceController.java b/src/com/android/settings/display/NightModePreferenceController.java
index 7e9701b..96ef938 100644
--- a/src/com/android/settings/display/NightModePreferenceController.java
+++ b/src/com/android/settings/display/NightModePreferenceController.java
@@ -13,6 +13,8 @@
  */
 package com.android.settings.display;
 
+import static android.content.Context.UI_MODE_SERVICE;
+
 import android.app.UiModeManager;
 import android.content.Context;
 import android.support.v7.preference.ListPreference;
@@ -23,8 +25,6 @@
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settingslib.core.AbstractPreferenceController;
 
-import static android.content.Context.UI_MODE_SERVICE;
-
 public class NightModePreferenceController extends AbstractPreferenceController implements
         PreferenceControllerMixin, Preference.OnPreferenceChangeListener {
 
@@ -48,7 +48,7 @@
     @Override
     public void displayPreference(PreferenceScreen screen) {
         if (!isAvailable()) {
-            removePreference(screen, KEY_NIGHT_MODE);
+            setVisible(screen, KEY_NIGHT_MODE, false /* visible */);
             return;
         }
         ListPreference mNightModePreference = (ListPreference) screen.findPreference(
diff --git a/src/com/android/settings/network/AirplaneModePreferenceController.java b/src/com/android/settings/network/AirplaneModePreferenceController.java
index d2015df..0620f14 100644
--- a/src/com/android/settings/network/AirplaneModePreferenceController.java
+++ b/src/com/android/settings/network/AirplaneModePreferenceController.java
@@ -81,7 +81,7 @@
                         mMetricsFeatureProvider);
             }
         } else {
-            removePreference(screen, getPreferenceKey());
+            setVisible(screen, getPreferenceKey(), false /* visible */);
         }
     }
 
@@ -91,7 +91,7 @@
     }
 
     public static boolean isAvailable(Context context) {
-        return !context.getPackageManager().hasSystemFeature(PackageManager.FEATURE_TELEVISION);
+        return !context.getPackageManager().hasSystemFeature(PackageManager.FEATURE_LEANBACK);
     }
 
     @Override
diff --git a/src/com/android/settings/nfc/NfcPreferenceController.java b/src/com/android/settings/nfc/NfcPreferenceController.java
index a034777..76977d2 100644
--- a/src/com/android/settings/nfc/NfcPreferenceController.java
+++ b/src/com/android/settings/nfc/NfcPreferenceController.java
@@ -26,9 +26,9 @@
 import android.support.v7.preference.PreferenceScreen;
 
 import com.android.settings.core.PreferenceControllerMixin;
+import com.android.settingslib.RestrictedPreference;
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
-import com.android.settingslib.RestrictedPreference;
 import com.android.settingslib.core.lifecycle.events.OnPause;
 import com.android.settingslib.core.lifecycle.events.OnResume;
 
@@ -55,8 +55,8 @@
     @Override
     public void displayPreference(PreferenceScreen screen) {
         if (!isAvailable()) {
-            removePreference(screen, KEY_TOGGLE_NFC);
-            removePreference(screen, KEY_ANDROID_BEAM_SETTINGS);
+            setVisible(screen, KEY_TOGGLE_NFC, false /* visible */);
+            setVisible(screen, KEY_ANDROID_BEAM_SETTINGS, false /* visible */);
             mNfcEnabler = null;
             return;
         }
diff --git a/src/com/android/settings/notification/LockScreenNotificationPreferenceController.java b/src/com/android/settings/notification/LockScreenNotificationPreferenceController.java
index 9855751..dd7b0fb 100644
--- a/src/com/android/settings/notification/LockScreenNotificationPreferenceController.java
+++ b/src/com/android/settings/notification/LockScreenNotificationPreferenceController.java
@@ -98,8 +98,8 @@
             mLockscreenProfile = (RestrictedDropDownPreference) screen.findPreference(
                     mWorkSettingKey);
         } else {
-            removePreference(screen, mWorkSettingKey);
-            removePreference(screen, mWorkSettingCategoryKey);
+            setVisible(screen, mWorkSettingKey, false /* visible */);
+            setVisible(screen, mWorkSettingCategoryKey, false /* visible */);
         }
         mSettingObserver = new SettingObserver();
         initLockScreenNotificationPrefDisplay();
diff --git a/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java
index d399838..faaf7db 100644
--- a/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java
@@ -17,7 +17,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -27,7 +26,6 @@
 import android.content.pm.UserInfo;
 import android.os.UserManager;
 import android.provider.Settings.Global;
-import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 
 import com.android.settings.TestConfig;
@@ -70,13 +68,12 @@
         when(mUserManager.getUserInfo(anyInt())).thenReturn(mUserInfo);
         when(mUserInfo.isAdmin()).thenReturn(false);
         final RestrictedSwitchPreference preference = mock(RestrictedSwitchPreference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
         when(preference.getKey()).thenReturn(mController.getPreferenceKey());
+        when(mScreen.findPreference(preference.getKey())).thenReturn(preference);
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        verify(preference).setVisible(false);
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java
index b8ad837..4715f88 100644
--- a/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java
@@ -15,19 +15,23 @@
  */
 package com.android.settings.accounts;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Mockito.when;
+
 import android.app.Fragment;
 import android.content.Context;
 import android.content.DialogInterface;
 import android.content.pm.UserInfo;
 import android.os.UserManager;
+import android.support.v14.preference.SwitchPreference;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
-import android.support.v14.preference.SwitchPreference;
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.TestConfig;
 
-import java.util.ArrayList;
-import java.util.List;
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -36,14 +40,8 @@
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowApplication;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
+import java.util.ArrayList;
+import java.util.List;
 
 @RunWith(SettingsRobolectricTestRunner.class)
 @Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
@@ -55,9 +53,8 @@
     private UserManager mUserManager;
     @Mock(answer = RETURNS_DEEP_STUBS)
     private Fragment mFragment;
-    @Mock
-    private Preference mPreference;
 
+    private Preference mPreference;
     private Context mContext;
     private AutoSyncDataPreferenceController mController;
     private AutoSyncDataPreferenceController.ConfirmAutoSyncChangeFragment mConfirmSyncFragment;
@@ -71,9 +68,9 @@
         mController = new AutoSyncDataPreferenceController(mContext, mFragment);
         mConfirmSyncFragment = new AutoSyncDataPreferenceController.ConfirmAutoSyncChangeFragment();
         mConfirmSyncFragment.setTargetFragment(mFragment, 0);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(mPreference);
-        when(mPreference.getKey()).thenReturn(mController.getPreferenceKey());
+        mPreference = new Preference(mContext);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -82,7 +79,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -92,7 +89,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
@@ -105,7 +102,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
@@ -119,7 +116,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -135,5 +132,4 @@
         mConfirmSyncFragment.onClick(null, DialogInterface.BUTTON_NEGATIVE);
         assertThat(preference.isChecked()).isFalse();
     }
-
 }
diff --git a/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java
index ab462c0..597fbd4 100644
--- a/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java
@@ -15,17 +15,21 @@
  */
 package com.android.settings.accounts;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Mockito.when;
+
 import android.app.Fragment;
 import android.content.Context;
 import android.content.pm.UserInfo;
 import android.os.UserManager;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.TestConfig;
 
-import java.util.ArrayList;
-import java.util.List;
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -34,12 +38,8 @@
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowApplication;
 
-import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
+import java.util.ArrayList;
+import java.util.List;
 
 @RunWith(SettingsRobolectricTestRunner.class)
 @Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
@@ -51,10 +51,9 @@
     private UserManager mUserManager;
     @Mock(answer = RETURNS_DEEP_STUBS)
     private Fragment mFragment;
-    @Mock
-    private Preference mPreference;
 
     private Context mContext;
+    private Preference mPreference;
     private AutoSyncPersonalDataPreferenceController mController;
 
     @Before
@@ -64,9 +63,9 @@
         shadowContext.setSystemService(Context.USER_SERVICE, mUserManager);
         mContext = shadowContext.getApplicationContext();
         mController = new AutoSyncPersonalDataPreferenceController(mContext, mFragment);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(mPreference);
-        when(mPreference.getKey()).thenReturn(mController.getPreferenceKey());
+        mPreference = new Preference(mContext);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -75,7 +74,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -85,7 +84,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -98,11 +97,11 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
-    public void displayPref_prefAvaiable_shouldDisplay() {
+    public void displayPref_prefAvailable_shouldDisplay() {
         List<UserInfo> infos = new ArrayList<>();
         infos.add(new UserInfo(1, "user 1", 0));
         infos.add(new UserInfo(2, "user 2", 0));
@@ -112,7 +111,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
 }
diff --git a/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java
index 0f7de46..29a011e 100644
--- a/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java
@@ -15,31 +15,6 @@
  */
 package com.android.settings.accounts;
 
-import android.content.Context;
-import android.content.Intent;
-import android.content.pm.ResolveInfo;
-import android.content.pm.UserInfo;
-import android.os.UserManager;
-import android.support.v7.preference.Preference;
-import android.support.v7.preference.PreferenceScreen;
-
-import com.android.settings.R;
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.TestConfig;
-import com.android.settings.search.SearchIndexableRaw;
-import com.android.settings.testutils.shadow.ShadowAccountManager;
-import com.android.settings.testutils.shadow.ShadowContentResolver;
-
-import java.util.ArrayList;
-import java.util.List;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-import org.robolectric.annotation.Config;
-import org.robolectric.shadows.ShadowApplication;
-
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Answers.RETURNS_DEEP_STUBS;
 import static org.mockito.Matchers.any;
@@ -49,8 +24,35 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ResolveInfo;
+import android.content.pm.UserInfo;
+import android.os.UserManager;
+import android.support.v7.preference.Preference;
+import android.support.v7.preference.PreferenceScreen;
+
+import com.android.settings.R;
+import com.android.settings.TestConfig;
+import com.android.settings.search.SearchIndexableRaw;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
+import com.android.settings.testutils.shadow.ShadowAccountManager;
+import com.android.settings.testutils.shadow.ShadowContentResolver;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowApplication;
+
+import java.util.ArrayList;
+import java.util.List;
+
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class EmergencyInfoPreferenceControllerTest {
 
     @Mock(answer = RETURNS_DEEP_STUBS)
@@ -61,15 +63,19 @@
     private UserManager mUserManager;
 
     private EmergencyInfoPreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mController = new EmergencyInfoPreferenceController(mContext);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
-    public void updateRawDataToIndex_prefUnavaiable_shouldNotUpdate() {
+    public void updateRawDataToIndex_prefUnavailable_shouldNotUpdate() {
         final List<SearchIndexableRaw> data = new ArrayList<>();
         when(mContext.getPackageManager().queryIntentActivities(
                 any(Intent.class), anyInt()))
@@ -81,7 +87,7 @@
     }
 
     @Test
-    public void updateRawDataToIndex_prefAvaiable_shouldUpdate() {
+    public void updateRawDataToIndex_prefAvailable_shouldUpdate() {
         final List<SearchIndexableRaw> data = new ArrayList<>();
         final List<ResolveInfo> infos = new ArrayList<>();
         infos.add(new ResolveInfo());
@@ -95,23 +101,18 @@
     }
 
     @Test
-    public void displayPref_prefUnAvaiable_shouldNotDisplay() {
+    public void displayPref_prefUnAvailable_shouldNotDisplay() {
         when(mContext.getPackageManager().queryIntentActivities(
                 any(Intent.class), anyInt()))
                 .thenReturn(null);
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
-    public void displayPref_prefAvaiable_shouldDisplay() {
-        final List<SearchIndexableRaw> data = new ArrayList<>();
+    public void displayPref_prefAvailable_shouldDisplay() {
         final List<ResolveInfo> infos = new ArrayList<>();
         infos.add(new ResolveInfo());
         when(mContext.getPackageManager().queryIntentActivities(
diff --git a/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java
index c61dfc6..2d75b8a 100644
--- a/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java
@@ -16,17 +16,25 @@
 
 package com.android.settings.development;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.PackageManager;
-import android.provider.Settings;
 import android.os.UserManager;
+import android.provider.Settings;
 import android.support.v14.preference.SwitchPreference;
-import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -36,23 +44,12 @@
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowApplication;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 /**
  * deprecated in favor of {@link BugReportInPowerPreferenceControllerV2}
  */
 @Deprecated
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class BugReportInPowerPreferenceControllerTest {
 
     @Mock(answer = RETURNS_DEEP_STUBS)
@@ -82,12 +79,10 @@
     @Test
     public void displayPreference_hasDebugRestriction_shouldRemovePreference() {
         when(mUserManager.hasUserRestriction(anyString())).thenReturn(true);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(mPreference);
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -96,7 +91,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java
index d705610..dffa461 100644
--- a/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java
@@ -16,66 +16,62 @@
 
 package com.android.settings.development;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.os.UserManager;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
-import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyBoolean;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 /**
  * deprecated in favor of {@link BugReportPreferenceControllerV2}
  */
 @Deprecated
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class BugReportPreferenceControllerTest {
 
     @Mock
     private Context mContext;
-    @Mock
-    private Preference mPreference;
     @Mock(answer = RETURNS_DEEP_STUBS)
     private PreferenceScreen mScreen;
     @Mock
     private UserManager mUserManager;
 
     private BugReportPreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         when(mContext.getSystemService(Context.USER_SERVICE)).thenReturn(mUserManager);
-        when(mScreen.findPreference(anyString())).thenReturn(mPreference);
         mController = new BugReportPreferenceController(mContext);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
     public void displayPreference_hasDebugRestriction_shouldRemovePreference() {
         when(mUserManager.hasUserRestriction(anyString())).thenReturn(true);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(mPreference);
-        when(mPreference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -84,7 +80,7 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
@@ -92,9 +88,10 @@
         when(mUserManager.hasUserRestriction(anyString())).thenReturn(true);
         mController.displayPreference(mScreen);
 
+        mPreference.setEnabled(false);
         mController.enablePreference(true);
 
-        verify(mPreference, never()).setEnabled(anyBoolean());
+        assertThat(mPreference.isEnabled()).isFalse();
     }
 
     @Test
@@ -102,9 +99,10 @@
         when(mUserManager.hasUserRestriction(anyString())).thenReturn(false);
         mController.displayPreference(mScreen);
 
+        mPreference.setEnabled(false);
         mController.enablePreference(true);
 
-        verify(mPreference).setEnabled(anyBoolean());
+        assertThat(mPreference.isEnabled()).isTrue();
     }
 
 }
diff --git a/tests/robotests/src/com/android/settings/development/VerifyAppsOverUsbPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/VerifyAppsOverUsbPreferenceControllerTest.java
index cea3d89..3208810 100644
--- a/tests/robotests/src/com/android/settings/development/VerifyAppsOverUsbPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/VerifyAppsOverUsbPreferenceControllerTest.java
@@ -30,8 +30,8 @@
 import android.provider.Settings.Global;
 import android.support.v7.preference.PreferenceScreen;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settingslib.RestrictedLockUtils.EnforcedAdmin;
 import com.android.settingslib.RestrictedSwitchPreference;
 
@@ -76,6 +76,7 @@
             return this;
         }
     }
+
     private final GlobalSetter mGlobals = new GlobalSetter();
 
     @Before
@@ -175,7 +176,9 @@
         when(mPreference.getKey()).thenReturn(mController.getPreferenceKey());
         when(mScreen.getPreferenceCount()).thenReturn(1);
         when(mScreen.getPreference(anyInt())).thenReturn(mPreference);
+
         mController.displayPreference(mScreen);
-        verify(mScreen).removePreference(mPreference);
+
+        verify(mPreference).setVisible(false);
     }
 }
\ No newline at end of file
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java
index 5c62220..32bcd60 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java
@@ -15,6 +15,10 @@
  */
 package com.android.settings.deviceinfo;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.os.Build;
 import android.os.UserManager;
@@ -22,8 +26,8 @@
 import android.support.v7.preference.PreferenceScreen;
 
 import com.android.settings.R;
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -36,30 +40,27 @@
 import java.util.ArrayList;
 import java.util.List;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class SystemUpdatePreferenceControllerTest {
 
     @Mock(answer = RETURNS_DEEP_STUBS)
     private Context mContext;
     @Mock
     private UserManager mUserManager;
-    @Mock(answer = RETURNS_DEEP_STUBS)
+    @Mock
     private PreferenceScreen mScreen;
 
     private SystemUpdatePreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mController = new SystemUpdatePreferenceController(mContext, mUserManager);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -86,41 +87,30 @@
 
     @Test
     public void displayPrefs_nothingAvailable_shouldNotDisplay() {
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
-
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
     public void updateState_shouldSetToAndroidVersion() {
-        final Preference preference = new Preference(RuntimeEnvironment.application);
         mController = new SystemUpdatePreferenceController(
                 RuntimeEnvironment.application, mUserManager);
-        mController.updateState(preference);
+        mController.updateState(mPreference);
 
-        assertThat(preference.getSummary())
+        assertThat(mPreference.getSummary())
                 .isEqualTo(RuntimeEnvironment.application.getString(R.string.about_summary,
                         Build.VERSION.RELEASE));
     }
 
     @Test
     public void displayPrefs_oneAvailable_shouldDisplayOne() {
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
-
         when(mContext.getResources().getBoolean(
                 R.bool.config_additional_system_update_setting_enable))
                 .thenReturn(true);
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 }
diff --git a/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java
index 61b6909..df438c1 100644
--- a/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java
@@ -16,14 +16,21 @@
 
 package com.android.settings.gestures;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.os.Bundle;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 import android.support.v7.preference.TwoStatePreference;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.widget.VideoPreference;
 import com.android.settingslib.core.lifecycle.Lifecycle;
 
@@ -33,19 +40,11 @@
 import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class GesturePreferenceControllerTest {
 
 
@@ -55,12 +54,17 @@
     private PreferenceScreen mScreen;
     @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     private Lifecycle mLifecycle;
+
     private TestPrefController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mController = new TestPrefController(mContext, mLifecycle);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -70,20 +74,16 @@
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
     public void display_configIsFalse_shouldNotDisplay() {
         mController.mIsPrefAvailable = false;
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java
index e4b3cf2..88a3bcb 100644
--- a/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java
@@ -16,14 +16,18 @@
 
 package com.android.settings.language;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.speech.tts.TextToSpeech;
 import android.speech.tts.TtsEngines;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -31,19 +35,14 @@
 import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
 import java.util.ArrayList;
 import java.util.List;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class TtsPreferenceControllerTest {
 
     @Mock(answer = Answers.RETURNS_DEEP_STUBS)
@@ -54,21 +53,28 @@
     private PreferenceScreen mScreen;
 
     private TtsPreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
         mController = new TtsPreferenceController(mContext, mTtsEngines);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
     public void testIsAvailable_ttsEngineEmpty_shouldReturnFalse() {
-
         // Not available when there is no engine.
         when(mTtsEngines.getEngines()).thenReturn(new ArrayList<>());
 
         assertThat(mController.isAvailable()).isFalse();
+
+        mController.displayPreference(mScreen);
+
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -78,21 +84,9 @@
         when(mTtsEngines.getEngines()).thenReturn(infolist);
 
         assertThat(mController.isAvailable()).isTrue();
-    }
-
-    @Test
-    public void displayPreference_notAvailable_shouldRemoveCategory() {
-        final Preference preference = mock(Preference.class);
-        final Preference category = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(2);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(mScreen.getPreference(1)).thenReturn(category);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
-        when(category.getKey()).thenReturn("voice_category");
 
         mController.displayPreference(mScreen);
 
-        // Remove preference.
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 }
diff --git a/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java
index 602bbd1..6041007 100644
--- a/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java
@@ -1,46 +1,46 @@
 package com.android.settings.location;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.provider.Settings;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
-import static org.mockito.Answers.RETURNS_DEEP_STUBS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class AppLocationPermissionPreferenceControllerTest {
-    @Mock
-    private Preference mPreference;
+
     @Mock(answer = RETURNS_DEEP_STUBS)
     private PreferenceScreen mScreen;
 
     private AppLocationPermissionPreferenceController mController;
 
-    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+    @Mock
     private Context mContext;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        when(mScreen.findPreference(anyString())).thenReturn(mPreference);
+        mContext = RuntimeEnvironment.application;
         mController = new AppLocationPermissionPreferenceController(mContext);
+        mPreference = new Preference(mContext);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -48,13 +48,10 @@
         Settings.System.putInt(mContext.getContentResolver(),
                 android.provider.Settings.Global.LOCATION_SETTINGS_LINK_TO_PERMISSIONS_ENABLED,
                 0);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(mPreference);
-        when(mPreference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
@@ -64,7 +61,6 @@
                 1);
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
-
 }
diff --git a/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java
index ac158b6..d294122 100644
--- a/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/BadgingNotificationPreferenceControllerTest.java
@@ -18,9 +18,7 @@
 
 import static android.provider.Settings.Secure.NOTIFICATION_BADGING;
 import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -43,11 +41,12 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowApplication;
 
 @RunWith(RobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class BadgingNotificationPreferenceControllerTest {
 
     @Mock(answer = Answers.RETURNS_DEEP_STUBS)
@@ -56,11 +55,15 @@
     private PreferenceScreen mScreen;
 
     private BadgingNotificationPreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mController = new BadgingNotificationPreferenceController(mContext);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -70,7 +73,7 @@
                 .thenReturn(true);
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
@@ -78,14 +81,10 @@
         when(mContext.getResources().
                 getBoolean(com.android.internal.R.bool.config_notificationBadging))
                 .thenReturn(false);
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java
index f6c6bf2..4ee5f5a 100644
--- a/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java
@@ -16,6 +16,12 @@
 
 package com.android.settings.notification;
 
+import static android.provider.Settings.System.NOTIFICATION_LIGHT_PULSE;
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.provider.Settings;
 import android.support.v7.preference.Preference;
@@ -31,18 +37,12 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowApplication;
 
-import static android.provider.Settings.System.NOTIFICATION_LIGHT_PULSE;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(RobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class PulseNotificationPreferenceControllerTest {
 
     @Mock(answer = Answers.RETURNS_DEEP_STUBS)
@@ -51,11 +51,15 @@
     private PreferenceScreen mScreen;
 
     private PulseNotificationPreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mController = new PulseNotificationPreferenceController(mContext);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
@@ -65,7 +69,7 @@
                 .thenReturn(true);
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
@@ -73,14 +77,10 @@
         when(mContext.getResources().
                 getBoolean(com.android.internal.R.bool.config_intrusiveNotificationLed))
                 .thenReturn(false);
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java
index 054309c..19cf649 100644
--- a/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java
@@ -16,6 +16,12 @@
 
 package com.android.settings.notification;
 
+import static android.provider.Settings.System.VIBRATE_WHEN_RINGING;
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.provider.Settings;
 import android.support.v7.preference.Preference;
@@ -23,26 +29,20 @@
 import android.support.v7.preference.TwoStatePreference;
 import android.telephony.TelephonyManager;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowApplication;
 
-import static android.provider.Settings.System.VIBRATE_WHEN_RINGING;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class VibrateWhenRingPreferenceControllerTest {
 
     @Mock
@@ -53,36 +53,34 @@
     private TelephonyManager mTelephonyManager;
 
     private VibrateWhenRingPreferenceController mController;
+    private Preference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         when(mContext.getSystemService(Context.TELEPHONY_SERVICE)).thenReturn(mTelephonyManager);
         mController = new VibrateWhenRingPreferenceController(mContext);
+        mPreference = new Preference(RuntimeEnvironment.application);
+        mPreference.setKey(mController.getPreferenceKey());
+        when(mScreen.findPreference(mPreference.getKey())).thenReturn(mPreference);
     }
 
     @Test
     public void display_voiceCapable_shouldDisplay() {
         when(mTelephonyManager.isVoiceCapable()).thenReturn(true);
-        when(mScreen.findPreference(mController.getPreferenceKey()))
-            .thenReturn(mock(Preference.class));
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen, never()).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isTrue();
     }
 
     @Test
     public void display_notVoiceCapable_shouldNotDisplay() {
         when(mTelephonyManager.isVoiceCapable()).thenReturn(false);
-        final Preference preference = mock(Preference.class);
-        when(mScreen.getPreferenceCount()).thenReturn(1);
-        when(mScreen.getPreference(0)).thenReturn(preference);
-        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mScreen);
 
-        verify(mScreen).removePreference(any(Preference.class));
+        assertThat(mPreference.isVisible()).isFalse();
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/wifi/p2p/P2pCategoryPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/wifi/p2p/P2pCategoryPreferenceControllerTest.java
index dd2ef52..e53fb87 100644
--- a/tests/robotests/src/com/android/settings/wifi/p2p/P2pCategoryPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/wifi/p2p/P2pCategoryPreferenceControllerTest.java
@@ -16,12 +16,19 @@
 
 package com.android.settings.wifi.p2p;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceCategory;
 import android.support.v7.preference.PreferenceScreen;
 
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
+import com.android.settings.testutils.SettingsRobolectricTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -31,13 +38,8 @@
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(SettingsRobolectricTestRunner.class)
-@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION_O)
 public class P2pCategoryPreferenceControllerTest {
 
     @Mock
@@ -81,7 +83,8 @@
         mController.addChild(pref);
 
         verify(mCategory).addPreference(pref);
-        verify(mCategory).setVisible(true);
+        verify(mCategory, atLeastOnce()).setVisible(true);
+        verify(mCategory, never()).setVisible(false);
     }
 
     @Test