diff --git a/packages/SystemUI/src/com/android/systemui/controls/ControlsServiceInfo.kt b/packages/SystemUI/src/com/android/systemui/controls/ControlsServiceInfo.kt
index 860149d..d3174f1 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/ControlsServiceInfo.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/ControlsServiceInfo.kt
@@ -26,9 +26,10 @@
 import android.content.pm.PackageManager.MATCH_DIRECT_BOOT_UNAWARE
 import android.content.pm.ResolveInfo
 import android.content.pm.ServiceInfo
+import android.graphics.drawable.Drawable
 import android.os.UserHandle
 import android.service.controls.ControlsProviderService
-import androidx.annotation.VisibleForTesting
+import android.util.IconDrawableFactory
 import androidx.annotation.WorkerThread
 import com.android.settingslib.applications.DefaultAppInfo
 import com.android.systemui.R
@@ -47,7 +48,7 @@
 
     init {
         val metadata = serviceInfo.metaData
-                ?.getString(ControlsProviderService.META_DATA_PANEL_ACTIVITY) ?: ""
+            ?.getString(ControlsProviderService.META_DATA_PANEL_ACTIVITY) ?: ""
         val unflatenned = ComponentName.unflattenFromString(metadata)
         if (unflatenned != null && unflatenned.packageName == componentName.packageName) {
             _panelActivity = unflatenned
@@ -71,21 +72,21 @@
 
     @WorkerThread
     fun resolvePanelActivity(
-            allowAllApps: Boolean = false
+        allowAllApps: Boolean = false
     ) {
         if (resolved) return
         resolved = true
         val validPackages = context.resources
-                .getStringArray(R.array.config_controlsPreferredPackages)
+            .getStringArray(R.array.config_controlsPreferredPackages)
         if (componentName.packageName !in validPackages && !allowAllApps) return
         panelActivity = _panelActivity?.let {
             val resolveInfos = mPm.queryIntentActivitiesAsUser(
-                    Intent().setComponent(it),
-                    PackageManager.ResolveInfoFlags.of(
-                            MATCH_DIRECT_BOOT_AWARE.toLong() or
-                                    MATCH_DIRECT_BOOT_UNAWARE.toLong()
-                    ),
-                    UserHandle.of(userId)
+                Intent().setComponent(it),
+                PackageManager.ResolveInfoFlags.of(
+                    MATCH_DIRECT_BOOT_AWARE.toLong() or
+                            MATCH_DIRECT_BOOT_UNAWARE.toLong()
+                ),
+                UserHandle.of(userId)
             )
             if (resolveInfos.isNotEmpty() && verifyResolveInfo(resolveInfos[0])) {
                 it
@@ -118,6 +119,27 @@
         }
     }
 
+    @WorkerThread
+    override fun loadLabel(): CharSequence {
+        return componentName?.let {
+            val appInfo = mPm.getApplicationInfoAsUser(componentName.packageName, 0, userId)
+            appInfo.loadLabel(mPm)
+        }
+            ?: packageItemInfo?.loadLabel(mPm)
+            ?: throw IllegalArgumentException("Package info is missing")
+    }
+
+    @WorkerThread
+    override fun loadIcon(): Drawable {
+        val packageName =
+            componentName?.packageName
+                ?: packageItemInfo?.packageName
+                ?: throw IllegalArgumentException("Package info is missing")
+        val factory = IconDrawableFactory.newInstance(context)
+        val appInfo = mPm.getApplicationInfoAsUser(packageName, 0, userId)
+        return factory.getBadgedIcon(appInfo)
+    }
+
     override fun equals(other: Any?): Boolean {
         return other is ControlsServiceInfo &&
                 userId == other.userId &&
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
index 605dc3f..1a620d2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
@@ -22,6 +22,7 @@
 import android.content.pm.ApplicationInfo
 import android.content.pm.PackageManager
 import android.content.pm.ServiceInfo
+import android.graphics.drawable.Drawable
 import android.os.UserHandle
 import android.service.controls.ControlsProviderService
 import android.testing.AndroidTestingRunner
@@ -69,6 +70,7 @@
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.clearInvocations
+import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.never
 import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
@@ -516,9 +518,9 @@
                 name = componentName.className
             }
         return spy(ControlsServiceInfo(mContext, serviceInfo)).apply {
-            `when`(loadLabel()).thenReturn(label)
-            `when`(loadIcon()).thenReturn(mock())
-            `when`(panelActivity).thenReturn(panelComponentName)
+            doReturn(label).whenever(this).loadLabel()
+            doReturn(mock<Drawable>()).whenever(this).loadIcon()
+            doReturn(panelComponentName).whenever(this).panelActivity
         }
     }