Fix existing DND Settings shortcuts to point to Modes

Fixes: 365545604
Test: atest ShortcutsUpdaterTest + manual (flag flip + reboot)
Flag: android.app.modes_ui
Change-Id: I28f7e3e69175e92611668fdfa655a817ffcc905e
diff --git a/src/com/android/settings/shortcut/ShortcutsUpdater.java b/src/com/android/settings/shortcut/ShortcutsUpdater.java
index 7479998..90a60fd 100644
--- a/src/com/android/settings/shortcut/ShortcutsUpdater.java
+++ b/src/com/android/settings/shortcut/ShortcutsUpdater.java
@@ -21,6 +21,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import android.app.Flags;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -29,6 +30,11 @@
 import android.content.pm.ShortcutInfo;
 import android.content.pm.ShortcutManager;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.settings.Settings;
+
 import java.util.ArrayList;
 import java.util.List;
 
@@ -43,23 +49,48 @@
      */
     public static void updatePinnedShortcuts(Context context) {
         ShortcutManager sm = checkNotNull(context.getSystemService(ShortcutManager.class));
-        PackageManager pm = context.getPackageManager();
 
         List<ShortcutInfo> updates = new ArrayList<>();
         for (ShortcutInfo info : sm.getPinnedShortcuts()) {
-            if (!info.getId().startsWith(SHORTCUT_ID_PREFIX)) {
-                continue;
+            ResolveInfo resolvedActivity = resolveActivity(context, info);
+            if (resolvedActivity != null) {
+                // Id is preserved to update an existing shortcut, but the activity it opens might
+                // be different, according to maybeGetReplacingComponent.
+                updates.add(Shortcuts.createShortcutInfo(context, info.getId(), resolvedActivity));
             }
-            ComponentName cn = ComponentName.unflattenFromString(
-                    info.getId().substring(SHORTCUT_ID_PREFIX.length()));
-            ResolveInfo ri = pm.resolveActivity(new Intent(SHORTCUT_PROBE).setComponent(cn), 0);
-            if (ri == null) {
-                continue;
-            }
-            updates.add(Shortcuts.createShortcutInfo(context, info.getId(), ri));
         }
         if (!updates.isEmpty()) {
             sm.updateShortcuts(updates);
         }
     }
+
+    @Nullable
+    private static ResolveInfo resolveActivity(Context context, ShortcutInfo shortcut) {
+        if (!shortcut.getId().startsWith(SHORTCUT_ID_PREFIX)) {
+            return null;
+        }
+
+        ComponentName cn = ComponentName.unflattenFromString(
+                shortcut.getId().substring(SHORTCUT_ID_PREFIX.length()));
+        if (cn == null) {
+            return null;
+        }
+
+        // Check if the componentName is obsolete and has been replaced by a different one.
+        cn = maybeGetReplacingComponent(context, cn);
+        PackageManager pm = context.getPackageManager();
+        return pm.resolveActivity(new Intent(SHORTCUT_PROBE).setComponent(cn), 0);
+    }
+
+    @NonNull
+    private static ComponentName maybeGetReplacingComponent(Context context, ComponentName cn) {
+        // ZenModeSettingsActivity is replaced by ModesSettingsActivity and will be deleted
+        // soon (so we shouldn't use ZenModeSettingsActivity.class).
+        if (Flags.modesApi() && Flags.modesUi()
+                && cn.getClassName().endsWith("Settings$ZenModeSettingsActivity")) {
+            return new ComponentName(context, Settings.ModesSettingsActivity.class);
+        }
+
+        return cn;
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/shortcut/ShortcutsUpdaterTest.java b/tests/robotests/src/com/android/settings/shortcut/ShortcutsUpdaterTest.java
index b157174..5324ff5 100644
--- a/tests/robotests/src/com/android/settings/shortcut/ShortcutsUpdaterTest.java
+++ b/tests/robotests/src/com/android/settings/shortcut/ShortcutsUpdaterTest.java
@@ -27,14 +27,19 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.Flags;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.ShortcutInfo;
 import android.content.pm.ShortcutManager;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 
 import com.android.settings.Settings;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -52,6 +57,9 @@
 
     private Context mContext;
 
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
     @Mock
     private ShortcutManager mShortcutManager;
     @Captor
@@ -60,14 +68,12 @@
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = spy(RuntimeEnvironment.application);
+        doReturn(mShortcutManager).when(mContext).getSystemService(eq(Context.SHORTCUT_SERVICE));
     }
 
     @Test
-    public void shortcutsUpdateTask() {
-        mContext = spy(RuntimeEnvironment.application);
-        doReturn(mShortcutManager).when(mContext).getSystemService(eq(Context.SHORTCUT_SERVICE));
-
+    public void updatePinnedShortcuts_updatesAllShortcuts() {
         final List<ShortcutInfo> pinnedShortcuts = Arrays.asList(
                 makeShortcut("d1"),
                 makeShortcut("d2"),
@@ -89,6 +95,50 @@
         assertThat(updates.get(1).getShortLabel().toString()).isEqualTo("Sound & vibration");
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    public void updatePinnedShortcuts_withModesFlag_replacesDndByModes() {
+        List<ShortcutInfo> shortcuts = List.of(
+                makeShortcut(Settings.ZenModeSettingsActivity.class));
+        when(mShortcutManager.getPinnedShortcuts()).thenReturn(shortcuts);
+
+        ShortcutsUpdater.updatePinnedShortcuts(mContext);
+
+        verify(mShortcutManager, times(1)).updateShortcuts(mListCaptor.capture());
+        final List<ShortcutInfo> updates = mListCaptor.getValue();
+        assertThat(updates).hasSize(1);
+
+        // Id hasn't changed, but intent and label has.
+        ComponentName zenCn = new ComponentName(mContext, Settings.ZenModeSettingsActivity.class);
+        ComponentName modesCn = new ComponentName(mContext, Settings.ModesSettingsActivity.class);
+        assertThat(updates.get(0).getId()).isEqualTo(
+                SHORTCUT_ID_PREFIX + zenCn.flattenToShortString());
+        assertThat(updates.get(0).getIntent().getComponent()).isEqualTo(modesCn);
+        assertThat(updates.get(0).getShortLabel().toString()).isEqualTo("Modes");
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    public void updatePinnedShortcuts_withoutModesFlag_leavesDndAlone() {
+        List<ShortcutInfo> shortcuts = List.of(
+                makeShortcut(Settings.ZenModeSettingsActivity.class));
+        when(mShortcutManager.getPinnedShortcuts()).thenReturn(shortcuts);
+
+        ShortcutsUpdater.updatePinnedShortcuts(mContext);
+
+        verify(mShortcutManager, times(1)).updateShortcuts(mListCaptor.capture());
+        final List<ShortcutInfo> updates = mListCaptor.getValue();
+        assertThat(updates).hasSize(1);
+
+        // Nothing has changed.
+        ComponentName zenCn = new ComponentName(mContext, Settings.ZenModeSettingsActivity.class);
+        assertThat(updates.get(0).getId()).isEqualTo(
+                SHORTCUT_ID_PREFIX + zenCn.flattenToShortString());
+        assertThat(updates.get(0).getIntent().getComponent()).isEqualTo(zenCn);
+        assertThat(updates.get(0).getShortLabel().toString()).isEqualTo("Do Not Disturb");
+
+    }
+
     private ShortcutInfo makeShortcut(Class<?> className) {
         ComponentName cn = new ComponentName(mContext, className);
         return makeShortcut(SHORTCUT_ID_PREFIX + cn.flattenToShortString());