diff --git a/src/com/android/settings/applications/managedomainurls/DomainAppPreference.java b/src/com/android/settings/applications/managedomainurls/DomainAppPreference.java
index e1f1e92cd0c..a9feedc49a6 100644
--- a/src/com/android/settings/applications/managedomainurls/DomainAppPreference.java
+++ b/src/com/android/settings/applications/managedomainurls/DomainAppPreference.java
@@ -20,52 +20,35 @@ import android.content.Context;
 import android.content.pm.PackageManager;
 import android.os.UserHandle;
 import android.util.ArraySet;
+import android.util.IconDrawableFactory;
 import android.view.View;
 
 import com.android.settings.R;
 import com.android.settings.Utils;
 import com.android.settings.widget.AppPreference;
-import com.android.settingslib.applications.ApplicationsState;
 import com.android.settingslib.applications.ApplicationsState.AppEntry;
 
-import androidx.annotation.VisibleForTesting;
 import androidx.preference.PreferenceViewHolder;
 
-@VisibleForTesting
 public class DomainAppPreference extends AppPreference {
+
     private final AppEntry mEntry;
     private final PackageManager mPm;
-    private final ApplicationsState mApplicationsState;
+    private final IconDrawableFactory mIconDrawableFactory;
 
-    public DomainAppPreference(final Context context, ApplicationsState applicationsState,
+    public DomainAppPreference(final Context context, IconDrawableFactory iconFactory,
             AppEntry entry) {
         super(context);
-        mApplicationsState = applicationsState;
+        mIconDrawableFactory = iconFactory;
         mPm = context.getPackageManager();
         mEntry = entry;
         mEntry.ensureLabel(getContext());
+
         setState();
-        if (mEntry.icon != null) {
-            setIcon(mEntry.icon);
-        }
     }
 
     @Override
     public void onBindViewHolder(PreferenceViewHolder holder) {
-        if (mEntry.icon == null) {
-            holder.itemView.post(new Runnable() {
-                @Override
-                public void run() {
-                    // Ensure we have an icon before binding.
-                    if (mApplicationsState != null) {
-                        mApplicationsState.ensureIcon(mEntry);
-                    }
-                    // This might trigger us to bind again, but it gives an easy way to only
-                    // load the icon once its needed, so its probably worth it.
-                    setIcon(mEntry.icon);
-                }
-            });
-        }
         super.onBindViewHolder(holder);
         holder.itemView.findViewById(R.id.appendix).setVisibility(View.GONE);
     }
@@ -81,6 +64,7 @@ public class DomainAppPreference extends AppPreference {
 
     private void setState() {
         setTitle(mEntry.label);
+        setIcon(mIconDrawableFactory.getBadgedIcon(mEntry.info));
         setSummary(getDomainsSummary(mEntry.info.packageName));
     }
 
diff --git a/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceController.java b/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceController.java
index 384a5cf1c00..a1b9ab5c136 100644
--- a/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceController.java
+++ b/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceController.java
@@ -20,6 +20,7 @@ import android.app.Application;
 import android.content.Context;
 import android.text.TextUtils;
 import android.util.ArrayMap;
+import android.util.IconDrawableFactory;
 
 import com.android.settings.R;
 import com.android.settings.applications.AppInfoBase;
@@ -154,12 +155,14 @@ public class DomainAppPreferenceController extends BasePreferenceController impl
     private void rebuildAppList(PreferenceGroup group, ArrayList<AppEntry> apps) {
         cacheAllPrefs(group);
         final int size = apps.size();
+        final Context context = group.getContext();
+        final IconDrawableFactory iconDrawableFactory = IconDrawableFactory.newInstance(context);
         for (int i = 0; i < size; i++) {
-            AppEntry entry = apps.get(i);
-            String key = entry.info.packageName + "|" + entry.info.uid;
+            final AppEntry entry = apps.get(i);
+            final String key = entry.info.packageName + "|" + entry.info.uid;
             DomainAppPreference preference = (DomainAppPreference) getCachedPreference(key);
             if (preference == null) {
-                preference = new DomainAppPreference(group.getContext(), mApplicationsState, entry);
+                preference = new DomainAppPreference(context, iconDrawableFactory, entry);
                 preference.setKey(key);
                 group.addPreference(preference);
             } else {
diff --git a/tests/robotests/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceControllerTest.java
index 312fb1f4bc4..2baacda8926 100644
--- a/tests/robotests/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/applications/managedomainurls/DomainAppPreferenceControllerTest.java
@@ -23,7 +23,7 @@ import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
-import android.graphics.drawable.Drawable;
+import android.util.IconDrawableFactory;
 import android.view.View;
 import android.widget.ProgressBar;
 
@@ -34,40 +34,39 @@ import com.android.settingslib.applications.ApplicationsState;
 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 java.util.UUID;
+
 import androidx.preference.PreferenceViewHolder;
 
 @RunWith(SettingsRobolectricTestRunner.class)
 public class DomainAppPreferenceControllerTest {
 
-    @Mock
     private ApplicationsState.AppEntry mAppEntry;
     private Context mContext;
+    private IconDrawableFactory mIconDrawableFactory;
 
     @Before
     public void setUp() {
-        MockitoAnnotations.initMocks(this);
         mContext = RuntimeEnvironment.application;
+        mIconDrawableFactory = IconDrawableFactory.newInstance(mContext);
+        mAppEntry = new ApplicationsState.AppEntry(
+                mContext, createApplicationInfo(mContext.getPackageName()), 0);
     }
 
     @Test
-    public void domainAppPreferenceShouldUseAppPreferenceLayout() {
-        mAppEntry.info = new ApplicationInfo();
-        mAppEntry.info.packageName = "com.android.settings.test";
-        final DomainAppPreference pref = new DomainAppPreference(mContext, null, mAppEntry);
+    public void getLayoutResource_shouldUseAppPreferenceLayout() {
+        final DomainAppPreference pref = new DomainAppPreference(
+                mContext, mIconDrawableFactory, mAppEntry);
 
         assertThat(pref.getLayoutResource()).isEqualTo(R.layout.preference_app);
     }
 
     @Test
     public void onBindViewHolder_shouldSetAppendixViewToGone() {
-        mAppEntry.info = new ApplicationInfo();
-        mAppEntry.info.packageName = "com.android.settings.test";
-        mAppEntry.icon = mock(Drawable.class);
-        final DomainAppPreference pref = new DomainAppPreference(mContext, null, mAppEntry);
+        final DomainAppPreference pref = new DomainAppPreference(
+                mContext, mIconDrawableFactory, mAppEntry);
         final View holderView = mock(View.class);
         final View appendixView = mock(View.class);
         when(holderView.findViewById(R.id.summary_container)).thenReturn(mock(View.class));
@@ -78,4 +77,13 @@ public class DomainAppPreferenceControllerTest {
 
         verify(appendixView).setVisibility(View.GONE);
     }
+
+    private ApplicationInfo createApplicationInfo(String packageName) {
+        ApplicationInfo appInfo = new ApplicationInfo();
+        appInfo.sourceDir = "foo";
+        appInfo.flags |= ApplicationInfo.FLAG_INSTALLED;
+        appInfo.storageUuid = UUID.randomUUID();
+        appInfo.packageName = packageName;
+        return appInfo;
+    }
 }