diff --git a/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java b/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java
index a5e83732ef2..41800a24ca8 100644
--- a/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java
+++ b/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java
@@ -97,7 +97,7 @@ public class SecondaryUserController extends PreferenceController implements
     @Override
     public void displayPreference(PreferenceScreen screen) {
         if (mStoragePreference == null) {
-            mStoragePreference = new StorageItemPreferenceAlternate(mContext);
+            mStoragePreference = new StorageItemPreferenceAlternate(screen.getContext());
 
             PreferenceGroup group =
                     (PreferenceGroup) screen.findPreference(TARGET_PREFERENCE_GROUP_KEY);
diff --git a/src/com/android/settings/deviceinfo/storage/UserProfileController.java b/src/com/android/settings/deviceinfo/storage/UserProfileController.java
index 5da9fec89ec..963784d6fdf 100644
--- a/src/com/android/settings/deviceinfo/storage/UserProfileController.java
+++ b/src/com/android/settings/deviceinfo/storage/UserProfileController.java
@@ -60,7 +60,7 @@ public class UserProfileController extends PreferenceController implements
 
     @Override
     public void displayPreference(PreferenceScreen screen) {
-        mStoragePreference = new StorageItemPreferenceAlternate(mContext);
+        mStoragePreference = new StorageItemPreferenceAlternate(screen.getContext());
         mStoragePreference.setOrder(mPreferenceOrder);
         mStoragePreference.setKey(PREFERENCE_KEY_BASE + mUser.id);
         mStoragePreference.setTitle(mUser.name);
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java
index 7222f535553..c8f3069f1d0 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java
@@ -55,6 +55,10 @@ public class SecondaryUserControllerTest {
     private static final String TARGET_PREFERENCE_GROUP_KEY = "pref_secondary_users";
     @Mock
     private UserManagerWrapper mUserManager;
+    @Mock
+    private PreferenceScreen mScreen;
+    @Mock
+    private PreferenceGroup mGroup;
 
     private Context mContext;
     private SecondaryUserController mController;
@@ -66,19 +70,19 @@ public class SecondaryUserControllerTest {
         mContext = RuntimeEnvironment.application;
         mPrimaryUser = new UserInfo();
         mController = new SecondaryUserController(mContext, mPrimaryUser);
+
+        when(mScreen.getContext()).thenReturn(mContext);
+        when(mScreen.findPreference(anyString())).thenReturn(mGroup);
+        when(mGroup.getKey()).thenReturn(TARGET_PREFERENCE_GROUP_KEY);
     }
 
     @Test
     public void controllerAddsSecondaryUser() throws Exception {
         mPrimaryUser.name = TEST_NAME;
-        PreferenceScreen screen = mock(PreferenceScreen.class);
-        PreferenceGroup group = mock(PreferenceGroup.class);
-        when(screen.findPreference(anyString())).thenReturn(group);
-        when(group.getKey()).thenReturn(TARGET_PREFERENCE_GROUP_KEY);
-        mController.displayPreference(screen);
+        mController.displayPreference(mScreen);
 
         final ArgumentCaptor<Preference> argumentCaptor = ArgumentCaptor.forClass(Preference.class);
-        verify(group).addPreference(argumentCaptor.capture());
+        verify(mGroup).addPreference(argumentCaptor.capture());
         Preference preference = argumentCaptor.getValue();
         assertThat(preference.getTitle()).isEqualTo(TEST_NAME);
     }
@@ -86,15 +90,11 @@ public class SecondaryUserControllerTest {
     @Test
     public void controllerUpdatesSummaryOfNewPreference() throws Exception {
         mPrimaryUser.name = TEST_NAME;
-        PreferenceScreen screen = mock(PreferenceScreen.class);
-        PreferenceGroup group = mock(PreferenceGroup.class);
-        when(screen.findPreference(anyString())).thenReturn(group);
-        when(group.getKey()).thenReturn(TARGET_PREFERENCE_GROUP_KEY);
-        mController.displayPreference(screen);
+        mController.displayPreference(mScreen);
         mController.setSize(10L);
         final ArgumentCaptor<Preference> argumentCaptor = ArgumentCaptor.forClass(Preference.class);
 
-        verify(group).addPreference(argumentCaptor.capture());
+        verify(mGroup).addPreference(argumentCaptor.capture());
 
         Preference preference = argumentCaptor.getValue();
         assertThat(preference.getSummary()).isEqualTo("10.00B");
@@ -153,11 +153,7 @@ public class SecondaryUserControllerTest {
     public void controllerUpdatesPreferenceOnAcceptingResult() throws Exception {
         mPrimaryUser.name = TEST_NAME;
         mPrimaryUser.id = 10;
-        PreferenceScreen screen = mock(PreferenceScreen.class);
-        PreferenceGroup group = mock(PreferenceGroup.class);
-        when(screen.findPreference(anyString())).thenReturn(group);
-        when(group.getKey()).thenReturn(TARGET_PREFERENCE_GROUP_KEY);
-        mController.displayPreference(screen);
+        mController.displayPreference(mScreen);
         StorageAsyncLoader.AppsStorageResult userResult =
                 new StorageAsyncLoader.AppsStorageResult();
         SparseArray<StorageAsyncLoader.AppsStorageResult> result = new SparseArray<>();
@@ -166,7 +162,7 @@ public class SecondaryUserControllerTest {
 
         mController.handleResult(result);
         final ArgumentCaptor<Preference> argumentCaptor = ArgumentCaptor.forClass(Preference.class);
-        verify(group).addPreference(argumentCaptor.capture());
+        verify(mGroup).addPreference(argumentCaptor.capture());
         Preference preference = argumentCaptor.getValue();
 
         assertThat(preference.getSummary()).isEqualTo("99.00B");
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java
index 2cd4f760118..ed49da48027 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java
@@ -21,6 +21,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.content.Intent;
@@ -53,6 +54,8 @@ public class UserProfileControllerTest {
 
     @Mock
     private UserManagerWrapper mUserManager;
+    @Mock
+    private PreferenceScreen mScreen;
 
     private Context mContext;
     private UserProfileController mController;
@@ -64,17 +67,17 @@ public class UserProfileControllerTest {
         mContext = spy(RuntimeEnvironment.application);
         mPrimaryProfile = new UserInfo();
         mController = new UserProfileController(mContext, mPrimaryProfile, 0);
+        when(mScreen.getContext()).thenReturn(mContext);
     }
 
     @Test
     public void controllerAddsPrimaryProfilePreference() throws Exception {
         mPrimaryProfile.name = TEST_NAME;
         mPrimaryProfile.id = 10;
-        PreferenceScreen screen = mock(PreferenceScreen.class);
-        mController.displayPreference(screen);
+        mController.displayPreference(mScreen);
 
         final ArgumentCaptor<Preference> argumentCaptor = ArgumentCaptor.forClass(Preference.class);
-        verify(screen).addPreference(argumentCaptor.capture());
+        verify(mScreen).addPreference(argumentCaptor.capture());
         Preference preference = argumentCaptor.getValue();
 
         assertThat(preference.getTitle()).isEqualTo(TEST_NAME);
@@ -85,11 +88,10 @@ public class UserProfileControllerTest {
     public void tappingProfilePreferenceSendsToStorageProfileFragment() throws Exception {
         mPrimaryProfile.name = TEST_NAME;
         mPrimaryProfile.id = 10;
-        PreferenceScreen screen = mock(PreferenceScreen.class);
-        mController.displayPreference(screen);
+        mController.displayPreference(mScreen);
 
         final ArgumentCaptor<Preference> argumentCaptor = ArgumentCaptor.forClass(Preference.class);
-        verify(screen).addPreference(argumentCaptor.capture());
+        verify(mScreen).addPreference(argumentCaptor.capture());
         Preference preference = argumentCaptor.getValue();
         assertThat(mController.handlePreferenceTreeClick(preference)).isTrue();
         final ArgumentCaptor<Intent> intentCaptor = ArgumentCaptor.forClass(Intent.class);
@@ -105,8 +107,7 @@ public class UserProfileControllerTest {
     public void acceptingResultUpdatesPreferenceSize() throws Exception {
         mPrimaryProfile.name = TEST_NAME;
         mPrimaryProfile.id = 10;
-        PreferenceScreen screen = mock(PreferenceScreen.class);
-        mController.displayPreference(screen);
+        mController.displayPreference(mScreen);
         SparseArray<StorageAsyncLoader.AppsStorageResult> result = new SparseArray<>();
         StorageAsyncLoader.AppsStorageResult userResult =
                 new StorageAsyncLoader.AppsStorageResult();
@@ -115,7 +116,7 @@ public class UserProfileControllerTest {
 
         mController.handleResult(result);
         final ArgumentCaptor<Preference> argumentCaptor = ArgumentCaptor.forClass(Preference.class);
-        verify(screen).addPreference(argumentCaptor.capture());
+        verify(mScreen).addPreference(argumentCaptor.capture());
         Preference preference = argumentCaptor.getValue();
 
         assertThat(preference.getSummary()).isEqualTo("99.00B");