Merge "Split ZenMode.getIcon() into getIconKey() and ZenIconLoader.getIcon()" into main

12 files changed, 365 insertions, 148 deletions

diff --git a/core/res/res/drawable/ic_zen_mode_type_special_dnd.xml b/core/res/res/drawable/ic_zen_mode_type_special_dnd.xml
new file mode 100644
index 000000000000..fee9b0e2f451
--- /dev/null
+++ b/core/res/res/drawable/ic_zen_mode_type_special_dnd.xml
@@ -0,0 +1,25 @@
+<!--
+Copyright (C) 2024 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:tint="?android:attr/colorControlNormal"
+    android:viewportHeight="960"
+    android:viewportWidth="960">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M280,520L680,520L680,440L280,440L280,520ZM480,880Q397,880 324,848.5Q251,817 197,763Q143,709 111.5,636Q80,563 80,480Q80,397 111.5,324Q143,251 197,197Q251,143 324,111.5Q397,80 480,80Q563,80 636,111.5Q709,143 763,197Q817,251 848.5,324Q880,397 880,480Q880,563 848.5,636Q817,709 763,763Q709,817 636,848.5Q563,880 480,880ZM480,800Q614,800 707,707Q800,614 800,480Q800,346 707,253Q614,160 480,160Q346,160 253,253Q160,346 160,480Q160,614 253,707Q346,800 480,800ZM480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Z" />
+</vector>
\ No newline at end of file
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index cf0fc6159f42..bbe661e78b1d 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5575,6 +5575,7 @@
   <java-symbol type="drawable" name="ic_zen_mode_type_schedule_time" />
   <java-symbol type="drawable" name="ic_zen_mode_type_theater" />
   <java-symbol type="drawable" name="ic_zen_mode_type_unknown" />
+  <java-symbol type="drawable" name="ic_zen_mode_type_special_dnd" />
 
   <!-- System notification for background user sound -->
   <java-symbol type="string" name="bg_user_sound_notification_title_alarm" />
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIcon.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIcon.java
new file mode 100644
index 000000000000..001701e309d8
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIcon.java
@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2024 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.settingslib.notification.modes;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import android.graphics.drawable.Drawable;
+
+import androidx.annotation.DrawableRes;
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+/**
+ * Icon of a Zen Mode, already loaded from the owner's resources (if specified) or from a default.
+ */
+public record ZenIcon(@NonNull Key key, @NonNull Drawable drawable) {
+
+    /**
+     * Key of a Zen Mode Icon.
+     *
+     * <p>{@link #resPackage()} will be null if the resource belongs to the system, and thus can
+     * be loaded with any {@code Context}.
+     */
+    public record Key(@Nullable String resPackage, @DrawableRes int resId) {
+
+        public Key {
+            checkArgument(resId != 0, "Resource id must be valid");
+        }
+
+        static Key forSystemResource(@DrawableRes int resId) {
+            return new Key(null, resId);
+        }
+    }
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java
new file mode 100644
index 000000000000..0a0b65b50e8b
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2024 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.settingslib.notification.modes;
+
+import android.app.AutomaticZenRule;
+
+import com.android.internal.R;
+
+import com.google.common.collect.ImmutableMap;
+
+/**
+ * Known icon keys for zen modes that lack a custom {@link AutomaticZenRule#getIconResId()}, based
+ * on their {@link ZenMode.Kind} and {@link ZenMode#getType}.
+ */
+class ZenIconKeys {
+
+    /** The icon for Do Not Disturb mode. */
+    static final ZenIcon.Key MANUAL_DND = ZenIcon.Key.forSystemResource(
+            R.drawable.ic_zen_mode_type_special_dnd);
+
+    /**
+     * The default icon for implicit modes (they can also have a specific icon, if the user has
+     * chosen one via Settings).
+     */
+    static final ZenIcon.Key IMPLICIT_MODE_DEFAULT = ZenIcon.Key.forSystemResource(
+            R.drawable.ic_zen_mode_type_unknown);
+
+    private static final ImmutableMap<Integer, ZenIcon.Key> TYPE_DEFAULTS = ImmutableMap.of(
+            AutomaticZenRule.TYPE_UNKNOWN,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_unknown),
+            AutomaticZenRule.TYPE_OTHER,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_other),
+            AutomaticZenRule.TYPE_SCHEDULE_TIME,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_schedule_time),
+            AutomaticZenRule.TYPE_SCHEDULE_CALENDAR,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_schedule_calendar),
+            AutomaticZenRule.TYPE_BEDTIME,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_bedtime),
+            AutomaticZenRule.TYPE_DRIVING,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_driving),
+            AutomaticZenRule.TYPE_IMMERSIVE,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_immersive),
+            AutomaticZenRule.TYPE_THEATER,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_theater),
+            AutomaticZenRule.TYPE_MANAGED,
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_managed)
+    );
+
+    private static final ZenIcon.Key FOR_UNEXPECTED_TYPE =
+            ZenIcon.Key.forSystemResource(R.drawable.ic_zen_mode_type_unknown);
+
+    /** Default icon descriptors per mode {@link AutomaticZenRule.Type}. */
+    static ZenIcon.Key forType(@AutomaticZenRule.Type int ruleType) {
+        return TYPE_DEFAULTS.getOrDefault(ruleType, FOR_UNEXPECTED_TYPE);
+    }
+
+    private ZenIconKeys() { }
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
index 271d5c49b903..fe0f98af1186 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
@@ -16,28 +16,24 @@
 
 package com.android.settingslib.notification.modes;
 
+import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 
-import static java.util.Objects.requireNonNull;
-
-import android.annotation.DrawableRes;
-import android.annotation.Nullable;
-import android.app.AutomaticZenRule;
 import android.content.Context;
 import android.graphics.drawable.AdaptiveIconDrawable;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.InsetDrawable;
-import android.service.notification.SystemZenRules;
 import android.text.TextUtils;
 import android.util.Log;
 import android.util.LruCache;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import com.google.common.util.concurrent.FluentFuture;
-import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -54,7 +50,7 @@ public class ZenIconLoader {
     @Nullable // Until first usage
     private static ZenIconLoader sInstance;
 
-    private final LruCache<String, Drawable> mCache;
+    private final LruCache<ZenIcon.Key, Drawable> mCache;
     private final ListeningExecutorService mBackgroundExecutor;
 
     public static ZenIconLoader getInstance() {
@@ -64,90 +60,85 @@ public class ZenIconLoader {
         return sInstance;
     }
 
+    /** Replaces the singleton instance of {@link ZenIconLoader} by the provided one. */
+    @VisibleForTesting(otherwise = VisibleForTesting.NONE)
+    public static void setInstance(@Nullable ZenIconLoader instance) {
+        sInstance = instance;
+    }
+
     private ZenIconLoader() {
         this(Executors.newFixedThreadPool(4));
     }
 
     @VisibleForTesting
-    ZenIconLoader(ExecutorService backgroundExecutor) {
+    public ZenIconLoader(ExecutorService backgroundExecutor) {
         mCache = new LruCache<>(50);
         mBackgroundExecutor =
                 MoreExecutors.listeningDecorator(backgroundExecutor);
     }
 
+    /**
+     * Loads the {@link Drawable} corresponding to a {@link ZenMode} in a background thread, and
+     * caches it for future calls.
+     *
+     * <p>The {@link ZenIcon#drawable()} will always correspond to the resource indicated by
+     * {@link ZenIcon#key()}. In turn, this will match the value of {@link ZenMode#getIconKey()}
+     * for the supplied mode -- except for the rare case where the mode has an apparently valid
+     * drawable resource id that we fail to load for some reason, thus needing a "fallback" icon.
+     */
     @NonNull
-    ListenableFuture<Drawable> getIcon(Context context, @NonNull AutomaticZenRule rule) {
-        if (rule.getIconResId() == 0) {
-            return Futures.immediateFuture(getFallbackIcon(context, rule.getType()));
-        }
+    public ListenableFuture<ZenIcon> getIcon(@NonNull Context context, @NonNull ZenMode mode) {
+        ZenIcon.Key key = mode.getIconKey();
+
+        return FluentFuture.from(loadIcon(context, key, /* useMonochromeIfPresent= */ true))
+                .transformAsync(drawable ->
+                        drawable != null
+                            ? immediateFuture(new ZenIcon(key, drawable))
+                            : getFallbackIcon(context, mode),
+                mBackgroundExecutor);
+    }
 
-        return FluentFuture.from(loadIcon(context, rule.getPackageName(), rule.getIconResId()))
-                .transform(icon ->
-                                icon != null ? icon : getFallbackIcon(context, rule.getType()),
-                        MoreExecutors.directExecutor());
+    private ListenableFuture<ZenIcon> getFallbackIcon(Context context, ZenMode mode) {
+        ZenIcon.Key key = ZenIconKeys.forType(mode.getType());
+        return FluentFuture.from(loadIcon(context, key, /* useMonochromeIfPresent= */ false))
+                .transform(drawable -> {
+                    checkNotNull(drawable, "Couldn't load DEFAULT icon for mode %s!", mode);
+                    return new ZenIcon(key, drawable);
+                },
+                directExecutor());
     }
 
     @NonNull
-    private ListenableFuture</* @Nullable */ Drawable> loadIcon(Context context, String pkg,
-            int iconResId) {
-        String cacheKey = pkg + ":" + iconResId;
+    private ListenableFuture</* @Nullable */ Drawable> loadIcon(Context context,
+            ZenIcon.Key key, boolean useMonochromeIfPresent) {
         synchronized (mCache) {
-            Drawable cachedValue = mCache.get(cacheKey);
+            Drawable cachedValue = mCache.get(key);
             if (cachedValue != null) {
                 return immediateFuture(cachedValue != MISSING ? cachedValue : null);
             }
         }
 
         return FluentFuture.from(mBackgroundExecutor.submit(() -> {
-            if (TextUtils.isEmpty(pkg) || SystemZenRules.PACKAGE_ANDROID.equals(pkg)) {
-                return context.getDrawable(iconResId);
+            if (TextUtils.isEmpty(key.resPackage())) {
+                return context.getDrawable(key.resId());
             } else {
-                Context appContext = context.createPackageContext(pkg, 0);
-                Drawable appDrawable = appContext.getDrawable(iconResId);
-                return getMonochromeIconIfPresent(appDrawable);
+                Context appContext = context.createPackageContext(key.resPackage(), 0);
+                Drawable appDrawable = appContext.getDrawable(key.resId());
+                return useMonochromeIfPresent
+                        ? getMonochromeIconIfPresent(appDrawable)
+                        : appDrawable;
             }
         })).catching(Exception.class, ex -> {
             // If we cannot resolve the icon, then store MISSING in the cache below, so
             // we don't try again.
-            Log.e(TAG, "Error while loading icon " + cacheKey, ex);
+            Log.e(TAG, "Error while loading mode icon " + key, ex);
             return null;
-        }, MoreExecutors.directExecutor()).transform(drawable -> {
+        }, directExecutor()).transform(drawable -> {
             synchronized (mCache) {
-                mCache.put(cacheKey, drawable != null ? drawable : MISSING);
+                mCache.put(key, drawable != null ? drawable : MISSING);
             }
             return drawable;
-        }, MoreExecutors.directExecutor());
-    }
-
-    private static Drawable getFallbackIcon(Context context, int ruleType) {
-        int iconResIdFromType = getIconResourceIdFromType(ruleType);
-        return requireNonNull(context.getDrawable(iconResIdFromType));
-    }
-
-    /** Return the default icon resource associated to a {@link AutomaticZenRule.Type} */
-    @DrawableRes
-    public static int getIconResourceIdFromType(@AutomaticZenRule.Type int ruleType) {
-        return switch (ruleType) {
-            case AutomaticZenRule.TYPE_UNKNOWN ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_unknown;
-            case AutomaticZenRule.TYPE_OTHER ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_other;
-            case AutomaticZenRule.TYPE_SCHEDULE_TIME ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_schedule_time;
-            case AutomaticZenRule.TYPE_SCHEDULE_CALENDAR ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_schedule_calendar;
-            case AutomaticZenRule.TYPE_BEDTIME ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_bedtime;
-            case AutomaticZenRule.TYPE_DRIVING ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_driving;
-            case AutomaticZenRule.TYPE_IMMERSIVE ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_immersive;
-            case AutomaticZenRule.TYPE_THEATER ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_theater;
-            case AutomaticZenRule.TYPE_MANAGED ->
-                    com.android.internal.R.drawable.ic_zen_mode_type_managed;
-            default -> com.android.internal.R.drawable.ic_zen_mode_type_unknown;
-        };
+        }, directExecutor());
     }
 
     private static Drawable getMonochromeIconIfPresent(Drawable icon) {
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
index 9fa8fc3cc647..36975c7ec4b1 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
@@ -20,9 +20,9 @@ import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_TIME;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_ALL;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
+import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
 import static android.service.notification.SystemZenRules.getTriggerDescriptionForScheduleEvent;
 import static android.service.notification.SystemZenRules.getTriggerDescriptionForScheduleTime;
-import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
 import static android.service.notification.ZenModeConfig.tryParseCountdownConditionId;
 import static android.service.notification.ZenModeConfig.tryParseEventConditionId;
 import static android.service.notification.ZenModeConfig.tryParseScheduleConditionId;
@@ -36,7 +36,6 @@ import android.annotation.SuppressLint;
 import android.app.AutomaticZenRule;
 import android.app.NotificationManager;
 import android.content.Context;
-import android.graphics.drawable.Drawable;
 import android.net.Uri;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -50,12 +49,8 @@ import androidx.annotation.DrawableRes;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.settingslib.R;
-
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
-import com.google.common.util.concurrent.Futures;
-import com.google.common.util.concurrent.ListenableFuture;
 
 import java.util.Comparator;
 import java.util.Objects;
@@ -275,28 +270,32 @@ public class ZenMode implements Parcelable {
     }
 
     /**
-     * Returns an icon "key" that is guaranteed to be different if the icon is different. Note that
-     * the inverse is not true, i.e. two keys can be different and the icon still be visually the
-     * same.
-     */
-    @NonNull
-    public String getIconKey() {
-        return mRule.getType() + ":" + mRule.getPackageName() + ":" + mRule.getIconResId();
-    }
-
-    /**
-     * Returns the mode icon -- which can be either app-provided (via {@code addAutomaticZenRule}),
-     * user-chosen (via the icon picker in Settings), or a default icon based on the mode type.
+     * Returns the {@link ZenIcon.Key} corresponding to the icon resource for this mode. This can be
+     * either app-provided (via {@link AutomaticZenRule#setIconResId}, user-chosen (via the icon
+     * picker in Settings), or a default icon based on the mode {@link Kind} and {@link #getType}.
      */
     @NonNull
-    public ListenableFuture<Drawable> getIcon(@NonNull Context context,
-            @NonNull ZenIconLoader iconLoader) {
-        if (mKind == Kind.MANUAL_DND) {
-            return Futures.immediateFuture(requireNonNull(
-                    context.getDrawable(R.drawable.ic_do_not_disturb_on_24dp)));
+    public ZenIcon.Key getIconKey() {
+        if (isManualDnd()) {
+            return ZenIconKeys.MANUAL_DND;
+        }
+        if (mRule.getIconResId() != 0) {
+            if (isSystemOwned()) {
+                // System-owned rules can only have system icons.
+                return ZenIcon.Key.forSystemResource(mRule.getIconResId());
+            } else {
+                // Technically, the icon of an app-provided rule could be a system icon if the
+                // user chose one with the picker. However, we cannot know for sure.
+                return new ZenIcon.Key(mRule.getPackageName(), mRule.getIconResId());
+            }
+        } else {
+            // Using a default icon (which is always a system icon).
+            if (mKind == Kind.IMPLICIT) {
+                return ZenIconKeys.IMPLICIT_MODE_DEFAULT;
+            } else {
+                return ZenIconKeys.forType(getType());
+            }
         }
-
-        return iconLoader.getIcon(context, mRule);
     }
 
     @NonNull
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenIconLoaderTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenIconLoaderTest.java
index 20461e37a786..6eb5f5b55d05 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenIconLoaderTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenIconLoaderTest.java
@@ -16,17 +16,12 @@
 
 package com.android.settingslib.notification.modes;
 
-import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import android.app.AutomaticZenRule;
 import android.content.Context;
-import android.graphics.drawable.Drawable;
-import android.net.Uri;
-import android.service.notification.ZenPolicy;
+import android.service.notification.SystemZenRules;
 
-import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
 
 import org.junit.Before;
@@ -48,44 +43,73 @@ public class ZenIconLoaderTest {
     }
 
     @Test
-    public void getIcon_systemOwnedRuleWithIcon_loads() throws Exception {
-        AutomaticZenRule systemRule = newRuleBuilder()
-                .setPackage("android")
+    public void getIcon_systemOwnedModeWithIcon_loads() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
                 .setIconResId(android.R.drawable.ic_media_play)
                 .build();
 
-        ListenableFuture<Drawable> loadFuture = mLoader.getIcon(mContext, systemRule);
-        assertThat(loadFuture.isDone()).isTrue();
-        assertThat(loadFuture.get()).isNotNull();
+        ZenIcon icon = mLoader.getIcon(mContext, mode).get();
+
+        assertThat(icon.drawable()).isNotNull();
+        assertThat(icon.key().resPackage()).isNull();
+        assertThat(icon.key().resId()).isEqualTo(android.R.drawable.ic_media_play);
     }
 
     @Test
-    public void getIcon_ruleWithoutSpecificIcon_loadsFallback() throws Exception {
-        AutomaticZenRule rule = newRuleBuilder()
+    public void getIcon_modeWithoutSpecificIcon_loadsFallback() throws Exception {
+        ZenMode mode = new TestModeBuilder()
                 .setType(AutomaticZenRule.TYPE_DRIVING)
                 .setPackage("com.blah")
                 .build();
 
-        ListenableFuture<Drawable> loadFuture = mLoader.getIcon(mContext, rule);
-        assertThat(loadFuture.isDone()).isTrue();
-        assertThat(loadFuture.get()).isNotNull();
+        ZenIcon icon = mLoader.getIcon(mContext, mode).get();
+
+        assertThat(icon.drawable()).isNotNull();
+        assertThat(icon.key().resPackage()).isNull();
+        assertThat(icon.key().resId()).isEqualTo(
+                com.android.internal.R.drawable.ic_zen_mode_type_driving);
     }
 
     @Test
     public void getIcon_ruleWithAppIconWithLoadFailure_loadsFallback() throws Exception {
-        AutomaticZenRule rule = newRuleBuilder()
+        ZenMode mode = new TestModeBuilder()
                 .setType(AutomaticZenRule.TYPE_DRIVING)
                 .setPackage("com.blah")
                 .setIconResId(-123456)
                 .build();
 
-        ListenableFuture<Drawable> loadFuture = mLoader.getIcon(mContext, rule);
-        assertThat(loadFuture.get()).isNotNull();
+        ZenIcon icon = mLoader.getIcon(mContext, mode).get();
+
+        assertThat(icon.drawable()).isNotNull();
+        assertThat(icon.key().resPackage()).isNull();
+        assertThat(icon.key().resId()).isEqualTo(
+                com.android.internal.R.drawable.ic_zen_mode_type_driving);
     }
 
-    private static AutomaticZenRule.Builder newRuleBuilder() {
-        return new AutomaticZenRule.Builder("Driving", Uri.parse("drive"))
-                .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
-                .setZenPolicy(new ZenPolicy.Builder().build());
+    @Test
+    public void getIcon_cachesCustomIcons() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                .setIconResId(android.R.drawable.ic_media_play)
+                .build();
+
+        ZenIcon iconOne = mLoader.getIcon(mContext, mode).get();
+        ZenIcon iconTwo = mLoader.getIcon(mContext, mode).get();
+
+        assertThat(iconOne.drawable()).isSameInstanceAs(iconTwo.drawable());
+    }
+
+    @Test
+    public void getIcon_cachesDefaultIcons() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                .setType(AutomaticZenRule.TYPE_IMMERSIVE)
+                .build();
+
+        ZenIcon iconOne = mLoader.getIcon(mContext, mode).get();
+        ZenIcon iconTwo = mLoader.getIcon(mContext, mode).get();
+
+        assertThat(iconOne.drawable()).isSameInstanceAs(iconTwo.drawable());
     }
 }
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
index 32216fadfb0d..e64b0c6d8e74 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
@@ -20,6 +20,7 @@ import static android.app.AutomaticZenRule.TYPE_BEDTIME;
 import static android.app.AutomaticZenRule.TYPE_DRIVING;
 import static android.app.AutomaticZenRule.TYPE_IMMERSIVE;
 import static android.app.AutomaticZenRule.TYPE_OTHER;
+import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
 import static android.app.AutomaticZenRule.TYPE_THEATER;
 import static android.app.AutomaticZenRule.TYPE_UNKNOWN;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_ALARMS;
@@ -30,14 +31,7 @@ import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-
 import android.app.AutomaticZenRule;
-import android.graphics.drawable.Drawable;
 import android.net.Uri;
 import android.os.Parcel;
 import android.service.notification.Condition;
@@ -47,12 +41,9 @@ import android.service.notification.ZenPolicy;
 
 import com.android.internal.R;
 
-import com.google.common.util.concurrent.ListenableFuture;
-
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -289,37 +280,97 @@ public class ZenModeTest {
     }
 
     @Test
-    public void getIcon_normalMode_loadsIconNormally() {
-        ZenIconLoader iconLoader = mock(ZenIconLoader.class);
-        ZenMode mode = new ZenMode("id", ZEN_RULE, zenConfigRuleFor(ZEN_RULE, false));
+    public void getIconKey_normalModeWithCustomIcon_isCustomIcon() {
+        ZenMode mode = new TestModeBuilder()
+                .setType(TYPE_BEDTIME)
+                .setPackage("some.package")
+                .setIconResId(123)
+                .build();
+
+        ZenIcon.Key iconKey = mode.getIconKey();
+
+        assertThat(iconKey.resPackage()).isEqualTo("some.package");
+        assertThat(iconKey.resId()).isEqualTo(123);
+    }
+
+    @Test
+    public void getIconKey_systemOwnedModeWithCustomIcon_isCustomIcon() {
+        ZenMode mode = new TestModeBuilder()
+                .setType(TYPE_SCHEDULE_CALENDAR)
+                .setPackage(PACKAGE_ANDROID)
+                .setIconResId(123)
+                .build();
 
-        ListenableFuture<Drawable> unused = mode.getIcon(RuntimeEnvironment.getApplication(),
-                iconLoader);
+        ZenIcon.Key iconKey = mode.getIconKey();
 
-        verify(iconLoader).getIcon(any(), eq(ZEN_RULE));
+        assertThat(iconKey.resPackage()).isNull();
+        assertThat(iconKey.resId()).isEqualTo(123);
     }
 
     @Test
-    public void getIcon_manualDnd_returnsFixedIcon() {
-        ZenIconLoader iconLoader = mock(ZenIconLoader.class);
+    public void getIconKey_implicitModeWithCustomIcon_isCustomIcon() {
+        ZenMode mode = new TestModeBuilder()
+                .setId(ZenModeConfig.implicitRuleId("some.package"))
+                .setPackage("some.package")
+                .setIconResId(123)
+                .build();
+
+        ZenIcon.Key iconKey = mode.getIconKey();
 
-        ListenableFuture<Drawable> future = TestModeBuilder.MANUAL_DND_INACTIVE.getIcon(
-                RuntimeEnvironment.getApplication(), iconLoader);
+        assertThat(iconKey.resPackage()).isEqualTo("some.package");
+        assertThat(iconKey.resId()).isEqualTo(123);
+    }
+
+    @Test
+    public void getIconKey_manualDnd_isDndIcon() {
+        ZenIcon.Key iconKey = TestModeBuilder.MANUAL_DND_INACTIVE.getIconKey();
 
-        assertThat(future.isDone()).isTrue();
-        verify(iconLoader, never()).getIcon(any(), any());
+        assertThat(iconKey.resPackage()).isNull();
+        assertThat(iconKey.resId()).isEqualTo(
+                com.android.internal.R.drawable.ic_zen_mode_type_special_dnd);
     }
 
     @Test
-    public void getIcon_implicitMode_loadsIconNormally() {
-        ZenIconLoader iconLoader = mock(ZenIconLoader.class);
-        ZenMode mode = new ZenMode(IMPLICIT_RULE_ID, IMPLICIT_ZEN_RULE,
-                zenConfigRuleFor(IMPLICIT_ZEN_RULE, false));
+    public void getIconKey_normalModeWithoutCustomIcon_isModeTypeIcon() {
+        ZenMode mode = new TestModeBuilder()
+                .setType(TYPE_BEDTIME)
+                .setPackage("some.package")
+                .build();
+
+        ZenIcon.Key iconKey = mode.getIconKey();
+
+        assertThat(iconKey.resPackage()).isNull();
+        assertThat(iconKey.resId()).isEqualTo(
+                com.android.internal.R.drawable.ic_zen_mode_type_bedtime);
+    }
+
+    @Test
+    public void getIconKey_systemOwnedModeWithoutCustomIcon_isModeTypeIcon() {
+        ZenMode mode = new TestModeBuilder()
+                .setType(TYPE_SCHEDULE_CALENDAR)
+                .setPackage(PACKAGE_ANDROID)
+                .build();
+
+        ZenIcon.Key iconKey = mode.getIconKey();
+
+        assertThat(iconKey.resPackage()).isNull();
+        assertThat(iconKey.resId()).isEqualTo(
+                com.android.internal.R.drawable.ic_zen_mode_type_schedule_calendar);
+    }
+
+    @Test
+    public void getIconKey_implicitModeWithoutCustomIcon_isSpecialIcon() {
+        ZenMode mode = new TestModeBuilder()
+                .setId(ZenModeConfig.implicitRuleId("some.package"))
+                .setPackage("some_package")
+                .setType(TYPE_BEDTIME) // Type should be ignored.
+                .build();
 
-        ListenableFuture<Drawable> unused = mode.getIcon(RuntimeEnvironment.getApplication(),
-                iconLoader);
+        ZenIcon.Key iconKey = mode.getIconKey();
 
-        verify(iconLoader).getIcon(any(), eq(IMPLICIT_ZEN_RULE));
+        assertThat(iconKey.resPackage()).isNull();
+        assertThat(iconKey.resId()).isEqualTo(
+                com.android.internal.R.drawable.ic_zen_mode_type_unknown);
     }
 
     private static void assertUnparceledIsEqualToOriginal(String type, ZenMode original) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
index e0a53f8817bc..5925819f27a7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
@@ -24,6 +24,7 @@ import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.settingslib.notification.modes.ZenIconLoader
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.coroutines.collectValues
@@ -35,6 +36,7 @@ import com.android.systemui.statusbar.policy.data.repository.fakeZenModeReposito
 import com.android.systemui.statusbar.policy.domain.interactor.zenModeInteractor
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.toCollection
@@ -61,6 +63,7 @@ class ModesTileDataInteractorTest : SysuiTestCase() {
             addOverride(com.android.internal.R.drawable.ic_zen_mode_type_bedtime, BEDTIME_DRAWABLE)
             addOverride(com.android.internal.R.drawable.ic_zen_mode_type_driving, DRIVING_DRAWABLE)
         }
+        ZenIconLoader.setInstance(ZenIconLoader(MoreExecutors.newDirectExecutorService()))
     }
 
     @EnableFlags(Flags.FLAG_MODES_UI)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
index d2bc54e09944..33f379d020b7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
@@ -23,6 +23,7 @@ import android.provider.Settings
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.notification.modes.TestModeBuilder
+import com.android.settingslib.notification.modes.ZenIconLoader
 import com.android.settingslib.notification.modes.ZenMode
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -34,10 +35,12 @@ import com.android.systemui.statusbar.policy.ui.dialog.mockModesDialogDelegate
 import com.android.systemui.statusbar.policy.ui.dialog.mockModesDialogEventLogger
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.clearInvocations
@@ -64,6 +67,11 @@ class ModesDialogViewModelTest : SysuiTestCase() {
             mockDialogEventLogger
         )
 
+    @Before
+    fun setUp() {
+        ZenIconLoader.setInstance(ZenIconLoader(MoreExecutors.newDirectExecutorService()))
+    }
+
     @Test
     fun tiles_filtersOutUserDisabledModes() =
         testScope.runTest {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
index d1b5160fd490..05bd1a7676ae 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
@@ -44,7 +44,6 @@ import android.view.View;
 
 import androidx.lifecycle.Observer;
 
-import com.android.settingslib.notification.modes.ZenIconLoader;
 import com.android.settingslib.notification.modes.ZenMode;
 import com.android.systemui.Flags;
 import com.android.systemui.broadcast.BroadcastDispatcher;
@@ -404,13 +403,9 @@ public class PhoneStatusBarPolicy
         boolean visible = mode != null;
         if (visible) {
             // TODO: b/360399800 - Get the resource id, package, and cached drawable from the mode;
-            //  this is a shortcut for testing (there should be no direct dependency on
-            //  ZenIconLoader here).
-            String resPackage = mode.isSystemOwned() ? null : mode.getRule().getPackageName();
-            int iconResId = mode.getRule().getIconResId();
-            if (iconResId == 0) {
-                iconResId = ZenIconLoader.getIconResourceIdFromType(mode.getType());
-            }
+            //  this is a shortcut for testing.
+            String resPackage = mode.getIconKey().resPackage();
+            int iconResId = mode.getIconKey().resId();
 
             mIconController.setResourceIcon(mSlotZen, resPackage, iconResId,
                     /* preloadedIcon= */ null, mode.getName());
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
index d351da68e1c6..a67b47a9a0c9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
@@ -92,7 +92,7 @@ constructor(
     }
 
     suspend fun getModeIcon(mode: ZenMode): Icon {
-        return mode.getIcon(context, iconLoader).await().asIcon()
+        return iconLoader.getIcon(context, mode).await().drawable().asIcon()
     }
 
     /**