diff --git a/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java
index f49bb5fb472..755212d4ce4 100644
--- a/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/AddUserWhenLockedPreferenceControllerTest.java
@@ -69,6 +69,10 @@ public class AddUserWhenLockedPreferenceControllerTest {
     public void displayPref_NotAdmin_shouldNotDisplay() {
         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());
 
         mController.displayPreference(mScreen);
 
diff --git a/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java
index 7fb6fd768f1..f7dd7a3b054 100644
--- a/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/AutoSyncDataPreferenceControllerTest.java
@@ -56,6 +56,8 @@ public class AutoSyncDataPreferenceControllerTest {
     private UserManager mUserManager;
     @Mock(answer = RETURNS_DEEP_STUBS)
     private Fragment mFragment;
+    @Mock
+    private Preference mPreference;
 
     private Context mContext;
     private AutoSyncDataPreferenceController mController;
@@ -70,6 +72,9 @@ public class AutoSyncDataPreferenceControllerTest {
         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());
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java
index 54ee658b51f..6ac0eaba7fd 100644
--- a/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/AutoSyncPersonalDataPreferenceControllerTest.java
@@ -52,6 +52,8 @@ public class AutoSyncPersonalDataPreferenceControllerTest {
     private UserManager mUserManager;
     @Mock(answer = RETURNS_DEEP_STUBS)
     private Fragment mFragment;
+    @Mock
+    private Preference mPreference;
 
     private Context mContext;
     private AutoSyncPersonalDataPreferenceController mController;
@@ -63,6 +65,9 @@ public class AutoSyncPersonalDataPreferenceControllerTest {
         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());
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java
index d31c0f63bec..10bb283696a 100644
--- a/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/EmergencyInfoPreferenceControllerTest.java
@@ -99,6 +99,10 @@ public class EmergencyInfoPreferenceControllerTest {
         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);
 
diff --git a/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java
index 4ca6609def7..902fb150690 100644
--- a/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/BugReportInPowerPreferenceControllerTest.java
@@ -80,6 +80,8 @@ public class BugReportInPowerPreferenceControllerTest {
     @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);
 
diff --git a/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java
index 377b4670d58..35ca5f8c656 100644
--- a/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/BugReportPreferenceControllerTest.java
@@ -68,6 +68,9 @@ public class BugReportPreferenceControllerTest {
     @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);
 
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/ManageStoragePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/ManageStoragePreferenceControllerTest.java
index 12b32183e74..14c6d72f9a5 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/ManageStoragePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/ManageStoragePreferenceControllerTest.java
@@ -34,6 +34,7 @@ import org.robolectric.annotation.Config;
 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.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -88,7 +89,12 @@ public class ManageStoragePreferenceControllerTest {
     }
 
     @Test
-    public void displayPref_prefAvaiable_shouldNotDisplay() {
+    public void displayPref_prefNotAvaiable_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));
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/SerialNumberPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/SerialNumberPreferenceControllerTest.java
index 1b87e6c3433..bf9c4120c79 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/SerialNumberPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/SerialNumberPreferenceControllerTest.java
@@ -71,9 +71,12 @@ public class SerialNumberPreferenceControllerTest {
 
     @Test
     public void testDisplay_noSerial_shouldHidePreference() {
-        when(mScreen.findPreference(anyString())).thenReturn(mock(Preference.class));
-
+        final Preference preference = mock(Preference.class);
+        when(mScreen.getPreferenceCount()).thenReturn(1);
+        when(mScreen.getPreference(0)).thenReturn(preference);
         mController = new SerialNumberPreferenceController(mContext, null);
+        when(preference.getKey()).thenReturn(mController.getPreferenceKey());
+
         mController.displayPreference(mScreen);
 
         verify(mScreen).removePreference(any(Preference.class));
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java
index 9ca218c2072..b35c60a5c9d 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/SystemUpdatePreferenceControllerTest.java
@@ -34,6 +34,7 @@ import org.robolectric.annotation.Config;
 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;
 
@@ -80,6 +81,11 @@ public class SystemUpdatePreferenceControllerTest {
 
     @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));
@@ -87,6 +93,11 @@ public class SystemUpdatePreferenceControllerTest {
 
     @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);
diff --git a/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java
index 01c192163c9..3dc6892b6ed 100644
--- a/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/gestures/GesturePreferenceControllerTest.java
@@ -73,8 +73,10 @@ public class GesturePreferenceControllerTest {
     @Test
     public void display_configIsFalse_shouldNotDisplay() {
         mController.mIsPrefAvailable = false;
-        when(mScreen.findPreference(mController.getPreferenceKey()))
-                .thenReturn(mock(Preference.class));
+        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);
 
@@ -84,7 +86,7 @@ public class GesturePreferenceControllerTest {
     @Test
     public void onStart_shouldStartVideoPreference() {
         final VideoPreference videoPreference = mock(VideoPreference.class);
-        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(videoPreference);
+        when(mScreen.findPreference(mController.getVideoPrefKey())).thenReturn(videoPreference);
         mController.mIsPrefAvailable = true;
 
         mController.displayPreference(mScreen);
@@ -96,7 +98,7 @@ public class GesturePreferenceControllerTest {
     @Test
     public void onStop_shouldStopVideoPreference() {
         final VideoPreference videoPreference = mock(VideoPreference.class);
-        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(videoPreference);
+        when(mScreen.findPreference(mController.getVideoPrefKey())).thenReturn(videoPreference);
         mController.mIsPrefAvailable = true;
 
         mController.displayPreference(mScreen);
@@ -163,12 +165,12 @@ public class GesturePreferenceControllerTest {
 
         @Override
         public String getPreferenceKey() {
-            return null;
+            return "testKey";
         }
 
         @Override
         protected String getVideoPrefKey() {
-            return null;
+            return "videoKey";
         }
 
         @Override
diff --git a/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java
index 0e4178413e1..31a3f07ce9f 100644
--- a/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/language/TtsPreferenceControllerTest.java
@@ -84,7 +84,13 @@ public class TtsPreferenceControllerTest {
 
     @Test
     public void displayPreference_notAvailable_shouldRemoveCategory() {
-        when(mScreen.findPreference(anyString())).thenReturn(mock(Preference.class));
+        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);
 
diff --git a/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java
index 8a035e5140e..c7f76df3c9f 100644
--- a/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/location/AppLocationPermissionPreferenceControllerTest.java
@@ -50,6 +50,9 @@ public class AppLocationPermissionPreferenceControllerTest {
         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);
 
diff --git a/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java
index 4eb33e989e5..f6c6bf2da4c 100644
--- a/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/PulseNotificationPreferenceControllerTest.java
@@ -73,8 +73,10 @@ public class PulseNotificationPreferenceControllerTest {
         when(mContext.getResources().
                 getBoolean(com.android.internal.R.bool.config_intrusiveNotificationLed))
                 .thenReturn(false);
-        when(mScreen.findPreference(mController.getPreferenceKey()))
-                .thenReturn(mock(Preference.class));
+        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);
 
diff --git a/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java
index 3c3f212087d..3bc1ffa2bdd 100644
--- a/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/VibrateWhenRingPreferenceControllerTest.java
@@ -75,8 +75,10 @@ public class VibrateWhenRingPreferenceControllerTest {
     @Test
     public void display_notVoiceCapable_shouldNotDisplay() {
         when(mTelephonyManager.isVoiceCapable()).thenReturn(false);
-        when(mScreen.findPreference(mController.getPreferenceKey()))
-                .thenReturn(mock(Preference.class));
+        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);