diff --git a/src/com/android/settings/widget/EntityHeaderController.java b/src/com/android/settings/widget/EntityHeaderController.java
index 6d67eb49039..5bf0e4c2df6 100644
--- a/src/com/android/settings/widget/EntityHeaderController.java
+++ b/src/com/android/settings/widget/EntityHeaderController.java
@@ -35,6 +35,7 @@ import android.support.annotation.IntDef;
 import android.support.annotation.VisibleForTesting;
 import android.support.v7.widget.RecyclerView;
 import android.text.TextUtils;
+import android.util.IconDrawableFactory;
 import android.util.Log;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -146,9 +147,7 @@ public class EntityHeaderController {
      * accessibility purposes.
      */
     public EntityHeaderController setIcon(ApplicationsState.AppEntry appEntry) {
-        if (appEntry.icon != null) {
-            mIcon = appEntry.icon.getConstantState().newDrawable(mAppContext.getResources());
-        }
+        mIcon = IconDrawableFactory.newInstance(mAppContext).getBadgedIcon(appEntry.info);
         return this;
     }
 
diff --git a/tests/robotests/src/com/android/settings/testutils/ShadowIconDrawableFactory.java b/tests/robotests/src/com/android/settings/testutils/ShadowIconDrawableFactory.java
new file mode 100644
index 00000000000..cfba1e1988c
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/testutils/ShadowIconDrawableFactory.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.testutils;
+
+import android.content.pm.ApplicationInfo;
+import android.graphics.Color;
+import android.graphics.drawable.ColorDrawable;
+import android.graphics.drawable.Drawable;
+import android.util.IconDrawableFactory;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+
+@Implements(IconDrawableFactory.class)
+public class ShadowIconDrawableFactory {
+
+    @Implementation
+    public Drawable getBadgedIcon(ApplicationInfo appInfo) {
+        return new ColorDrawable(Color.BLUE);
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/widget/EntityHeaderControllerTest.java b/tests/robotests/src/com/android/settings/widget/EntityHeaderControllerTest.java
index 22ad764ef9c..cdc17ae9799 100644
--- a/tests/robotests/src/com/android/settings/widget/EntityHeaderControllerTest.java
+++ b/tests/robotests/src/com/android/settings/widget/EntityHeaderControllerTest.java
@@ -31,6 +31,7 @@ import android.app.Fragment;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
+import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.ResolveInfo;
 import android.graphics.drawable.ColorDrawable;
@@ -39,12 +40,14 @@ import android.support.v7.preference.Preference;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.widget.ImageButton;
+import android.widget.ImageView;
 import android.widget.TextView;
 
 import com.android.settings.R;
 import com.android.settings.applications.LayoutPreference;
-import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
+import com.android.settings.testutils.ShadowIconDrawableFactory;
+import com.android.settingslib.applications.ApplicationsState;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -53,6 +56,7 @@ import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
 
 @RunWith(SettingsRobolectricTestRunner.class)
 public class EntityHeaderControllerTest {
@@ -68,12 +72,10 @@ public class EntityHeaderControllerTest {
     private LayoutInflater mLayoutInflater;
     private PackageInfo mInfo;
     private EntityHeaderController mController;
-    private FakeFeatureFactory mFeatureFactory;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mFeatureFactory = FakeFeatureFactory.setupForTest();
         mShadowContext = RuntimeEnvironment.application;
         when(mActivity.getApplicationContext()).thenReturn(mShadowContext);
         when(mContext.getApplicationContext()).thenReturn(mContext);
@@ -254,6 +256,23 @@ public class EntityHeaderControllerTest {
                 .isEqualTo(description);
     }
 
+    @Test
+    @Config(shadows = ShadowIconDrawableFactory.class)
+    public void setIcon_usingAppEntry_shouldLoadIconFromDrawableFactory() {
+        final View view = mLayoutInflater
+                .inflate(R.layout.settings_entity_header, null /* root */);
+        final ApplicationsState.AppEntry entry = mock(ApplicationsState.AppEntry.class);
+        entry.info = new ApplicationInfo();
+        mController = EntityHeaderController.newInstance(mActivity, mFragment, view);
+        mController.setIcon(entry).done(mActivity);
+        final ImageView iconView = view.findViewById(R.id.entity_header_icon);
+
+        // Icon is set
+        assertThat(iconView.getDrawable()).isNotNull();
+        // ... and entry.icon is still empty. This means the icon didn't come from cache.
+        assertThat(entry.icon).isNull();
+    }
+
     @Test
     public void bindButton_hasAppNotifIntent_shouldShowButton() {
         final View appLinks = mLayoutInflater