Merge "15/ Move some more logic into gesture/device state" into ub-launcher3-master
diff --git a/src/com/android/launcher3/notification/NotificationListener.java b/src/com/android/launcher3/notification/NotificationListener.java
index 10378ee..0e73829 100644
--- a/src/com/android/launcher3/notification/NotificationListener.java
+++ b/src/com/android/launcher3/notification/NotificationListener.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.notification;
 
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.SecureSettingsObserver.newNotificationSettingsObserver;
 
@@ -32,9 +33,10 @@
 import android.util.Log;
 import android.util.Pair;
 
+import androidx.annotation.AnyThread;
 import androidx.annotation.Nullable;
+import androidx.annotation.WorkerThread;
 
-import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.SecureSettingsObserver;
 
@@ -44,6 +46,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 /**
  * A {@link NotificationListenerService} that sends updates to its
@@ -59,16 +62,17 @@
     private static final int MSG_NOTIFICATION_POSTED = 1;
     private static final int MSG_NOTIFICATION_REMOVED = 2;
     private static final int MSG_NOTIFICATION_FULL_REFRESH = 3;
+    private static final int MSG_CANCEL_NOTIFICATION = 4;
+    private static final int MSG_RANKING_UPDATE = 5;
 
     private static NotificationListener sNotificationListenerInstance = null;
     private static NotificationsChangedListener sNotificationsChangedListener;
-    private static StatusBarNotificationsChangedListener sStatusBarNotificationsChangedListener;
     private static boolean sIsConnected;
-    private static boolean sIsCreated;
 
     private final Handler mWorkerHandler;
     private final Handler mUiHandler;
     private final Ranking mTempRanking = new Ranking();
+
     /** Maps groupKey's to the corresponding group of notifications. */
     private final Map<String, NotificationGroup> mNotificationGroupMap = new HashMap<>();
     /** Maps keys to their corresponding current group key */
@@ -79,85 +83,12 @@
 
     private SecureSettingsObserver mNotificationDotsObserver;
 
-    private final Handler.Callback mWorkerCallback = new Handler.Callback() {
-        @Override
-        public boolean handleMessage(Message message) {
-            switch (message.what) {
-                case MSG_NOTIFICATION_POSTED:
-                    mUiHandler.obtainMessage(message.what, message.obj).sendToTarget();
-                    break;
-                case MSG_NOTIFICATION_REMOVED:
-                    mUiHandler.obtainMessage(message.what, message.obj).sendToTarget();
-                    break;
-                case MSG_NOTIFICATION_FULL_REFRESH:
-                    List<StatusBarNotification> activeNotifications;
-                    if (sIsConnected) {
-                        try {
-                            activeNotifications = filterNotifications(getActiveNotifications());
-                        } catch (SecurityException ex) {
-                            Log.e(TAG, "SecurityException: failed to fetch notifications");
-                            activeNotifications = new ArrayList<StatusBarNotification>();
-
-                        }
-                    } else {
-                        activeNotifications = new ArrayList<StatusBarNotification>();
-                    }
-
-                    mUiHandler.obtainMessage(message.what, activeNotifications).sendToTarget();
-                    break;
-            }
-            return true;
-        }
-    };
-
-    private final Handler.Callback mUiCallback = new Handler.Callback() {
-        @Override
-        public boolean handleMessage(Message message) {
-            switch (message.what) {
-                case MSG_NOTIFICATION_POSTED:
-                    if (sNotificationsChangedListener != null) {
-                        NotificationPostedMsg msg = (NotificationPostedMsg) message.obj;
-                        sNotificationsChangedListener.onNotificationPosted(msg.packageUserKey,
-                                msg.notificationKey, msg.shouldBeFilteredOut);
-                    }
-                    break;
-                case MSG_NOTIFICATION_REMOVED:
-                    if (sNotificationsChangedListener != null) {
-                        Pair<PackageUserKey, NotificationKeyData> pair
-                                = (Pair<PackageUserKey, NotificationKeyData>) message.obj;
-                        sNotificationsChangedListener.onNotificationRemoved(pair.first, pair.second);
-                    }
-                    break;
-                case MSG_NOTIFICATION_FULL_REFRESH:
-                    if (sNotificationsChangedListener != null) {
-                        sNotificationsChangedListener.onNotificationFullRefresh(
-                                (List<StatusBarNotification>) message.obj);
-                    }
-                    break;
-            }
-            return true;
-        }
-    };
-
     public NotificationListener() {
-        super();
-        mWorkerHandler = new Handler(MODEL_EXECUTOR.getLooper(), mWorkerCallback);
-        mUiHandler = new Handler(Looper.getMainLooper(), mUiCallback);
+        mWorkerHandler = new Handler(MODEL_EXECUTOR.getLooper(), this::handleWorkerMessage);
+        mUiHandler = new Handler(Looper.getMainLooper(), this::handleUiMessage);
         sNotificationListenerInstance = this;
     }
 
-    @Override
-    public void onCreate() {
-        super.onCreate();
-        sIsCreated = true;
-    }
-
-    @Override
-    public void onDestroy() {
-        super.onDestroy();
-        sIsCreated = false;
-    }
-
     public static @Nullable NotificationListener getInstanceIfConnected() {
         return sIsConnected ? sNotificationListenerInstance : null;
     }
@@ -168,25 +99,107 @@
         NotificationListener notificationListener = getInstanceIfConnected();
         if (notificationListener != null) {
             notificationListener.onNotificationFullRefresh();
-        } else if (!sIsCreated && sNotificationsChangedListener != null) {
+        } else {
             // User turned off dots globally, so we unbound this service;
             // tell the listener that there are no notifications to remove dots.
-            sNotificationsChangedListener.onNotificationFullRefresh(
-                    Collections.<StatusBarNotification>emptyList());
+            MODEL_EXECUTOR.submit(() -> MAIN_EXECUTOR.submit(() ->
+                            listener.onNotificationFullRefresh(Collections.emptyList())));
         }
     }
 
-    public static void setStatusBarNotificationsChangedListener
-            (StatusBarNotificationsChangedListener listener) {
-        sStatusBarNotificationsChangedListener = listener;
-    }
-
     public static void removeNotificationsChangedListener() {
         sNotificationsChangedListener = null;
     }
 
-    public static void removeStatusBarNotificationsChangedListener() {
-        sStatusBarNotificationsChangedListener = null;
+    private boolean handleWorkerMessage(Message message) {
+        switch (message.what) {
+            case MSG_NOTIFICATION_POSTED: {
+                StatusBarNotification sbn = (StatusBarNotification) message.obj;
+                mUiHandler.obtainMessage(shouldBeFilteredOut(sbn)
+                                ? MSG_NOTIFICATION_REMOVED : MSG_NOTIFICATION_POSTED,
+                        toKeyPair(sbn)).sendToTarget();
+                return true;
+            }
+            case MSG_NOTIFICATION_REMOVED: {
+                StatusBarNotification sbn = (StatusBarNotification) message.obj;
+                mUiHandler.obtainMessage(MSG_NOTIFICATION_REMOVED,
+                        toKeyPair(sbn)).sendToTarget();
+
+                NotificationGroup notificationGroup = mNotificationGroupMap.get(sbn.getGroupKey());
+                String key = sbn.getKey();
+                if (notificationGroup != null) {
+                    notificationGroup.removeChildKey(key);
+                    if (notificationGroup.isEmpty()) {
+                        if (key.equals(mLastKeyDismissedByLauncher)) {
+                            // Only cancel the group notification if launcher dismissed the
+                            // last child.
+                            cancelNotification(notificationGroup.getGroupSummaryKey());
+                        }
+                        mNotificationGroupMap.remove(sbn.getGroupKey());
+                    }
+                }
+                if (key.equals(mLastKeyDismissedByLauncher)) {
+                    mLastKeyDismissedByLauncher = null;
+                }
+                return true;
+            }
+            case MSG_NOTIFICATION_FULL_REFRESH:
+                List<StatusBarNotification> activeNotifications = null;
+                if (sIsConnected) {
+                    try {
+                        activeNotifications = Arrays.stream(getActiveNotifications())
+                                .filter(this::shouldBeFilteredOut)
+                                .collect(Collectors.toList());
+                    } catch (SecurityException ex) {
+                        Log.e(TAG, "SecurityException: failed to fetch notifications");
+                        activeNotifications = new ArrayList<>();
+                    }
+                } else {
+                    activeNotifications = new ArrayList<>();
+                }
+
+                mUiHandler.obtainMessage(message.what, activeNotifications).sendToTarget();
+                return true;
+            case MSG_CANCEL_NOTIFICATION: {
+                mLastKeyDismissedByLauncher = (String) message.obj;
+                cancelNotification(mLastKeyDismissedByLauncher);
+                return true;
+            }
+            case MSG_RANKING_UPDATE: {
+                String[] keys = ((RankingMap) message.obj).getOrderedKeys();
+                for (StatusBarNotification sbn : getActiveNotifications(keys)) {
+                    updateGroupKeyIfNecessary(sbn);
+                }
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private boolean handleUiMessage(Message message) {
+        switch (message.what) {
+            case MSG_NOTIFICATION_POSTED:
+                if (sNotificationsChangedListener != null) {
+                    Pair<PackageUserKey, NotificationKeyData> msg = (Pair) message.obj;
+                    sNotificationsChangedListener.onNotificationPosted(
+                            msg.first, msg.second);
+                }
+                break;
+            case MSG_NOTIFICATION_REMOVED:
+                if (sNotificationsChangedListener != null) {
+                    Pair<PackageUserKey, NotificationKeyData> msg = (Pair) message.obj;
+                    sNotificationsChangedListener.onNotificationRemoved(
+                            msg.first, msg.second);
+                }
+                break;
+            case MSG_NOTIFICATION_FULL_REFRESH:
+                if (sNotificationsChangedListener != null) {
+                    sNotificationsChangedListener.onNotificationFullRefresh(
+                            (List<StatusBarNotification>) message.obj);
+                }
+                break;
+        }
+        return true;
     }
 
     @Override
@@ -217,84 +230,37 @@
         super.onListenerDisconnected();
         sIsConnected = false;
         mNotificationDotsObserver.unregister();
+        onNotificationFullRefresh();
     }
 
     @Override
     public void onNotificationPosted(final StatusBarNotification sbn) {
-        super.onNotificationPosted(sbn);
-        if (sbn == null) {
-            // There is a bug in platform where we can get a null notification; just ignore it.
-            return;
-        }
-        mWorkerHandler.obtainMessage(MSG_NOTIFICATION_POSTED, new NotificationPostedMsg(sbn))
-            .sendToTarget();
-        if (sStatusBarNotificationsChangedListener != null) {
-            sStatusBarNotificationsChangedListener.onNotificationPosted(sbn);
-        }
-    }
-
-    /**
-     * An object containing data to send to MSG_NOTIFICATION_POSTED targets.
-     */
-    private class NotificationPostedMsg {
-        final PackageUserKey packageUserKey;
-        final NotificationKeyData notificationKey;
-        final boolean shouldBeFilteredOut;
-
-        NotificationPostedMsg(StatusBarNotification sbn) {
-            packageUserKey = PackageUserKey.fromNotification(sbn);
-            notificationKey = NotificationKeyData.fromNotification(sbn);
-            shouldBeFilteredOut = shouldBeFilteredOut(sbn);
+        if (sbn != null) {
+            mWorkerHandler.obtainMessage(MSG_NOTIFICATION_POSTED, sbn).sendToTarget();
         }
     }
 
     @Override
     public void onNotificationRemoved(final StatusBarNotification sbn) {
-        super.onNotificationRemoved(sbn);
-        if (sbn == null) {
-            // There is a bug in platform where we can get a null notification; just ignore it.
-            return;
+        if (sbn != null) {
+            mWorkerHandler.obtainMessage(MSG_NOTIFICATION_REMOVED, sbn).sendToTarget();
         }
-        Pair<PackageUserKey, NotificationKeyData> packageUserKeyAndNotificationKey
-            = new Pair<>(PackageUserKey.fromNotification(sbn),
-            NotificationKeyData.fromNotification(sbn));
-        mWorkerHandler.obtainMessage(MSG_NOTIFICATION_REMOVED, packageUserKeyAndNotificationKey)
-            .sendToTarget();
-        if (sStatusBarNotificationsChangedListener != null) {
-            sStatusBarNotificationsChangedListener.onNotificationRemoved(sbn);
-        }
-
-        NotificationGroup notificationGroup = mNotificationGroupMap.get(sbn.getGroupKey());
-        String key = sbn.getKey();
-        if (notificationGroup != null) {
-            notificationGroup.removeChildKey(key);
-            if (notificationGroup.isEmpty()) {
-                if (key.equals(mLastKeyDismissedByLauncher)) {
-                    // Only cancel the group notification if launcher dismissed the last child.
-                    cancelNotification(notificationGroup.getGroupSummaryKey());
-                }
-                mNotificationGroupMap.remove(sbn.getGroupKey());
-            }
-        }
-        if (key.equals(mLastKeyDismissedByLauncher)) {
-            mLastKeyDismissedByLauncher = null;
-        }
-    }
-
-    public void cancelNotificationFromLauncher(String key) {
-        mLastKeyDismissedByLauncher = key;
-        cancelNotification(key);
     }
 
     @Override
     public void onNotificationRankingUpdate(RankingMap rankingMap) {
-        super.onNotificationRankingUpdate(rankingMap);
-        String[] keys = rankingMap.getOrderedKeys();
-        for (StatusBarNotification sbn : getActiveNotifications(keys)) {
-            updateGroupKeyIfNecessary(sbn);
-        }
+        mWorkerHandler.obtainMessage(MSG_RANKING_UPDATE, rankingMap).sendToTarget();
     }
 
+    /**
+     * Cancels a notification
+     */
+    @AnyThread
+    public void cancelNotificationFromLauncher(String key) {
+        mWorkerHandler.obtainMessage(MSG_CANCEL_NOTIFICATION, key).sendToTarget();
+    }
+
+    @WorkerThread
     private void updateGroupKeyIfNecessary(StatusBarNotification sbn) {
         String childKey = sbn.getKey();
         String oldGroupKey = mNotificationGroupKeyMap.get(childKey);
@@ -328,43 +294,23 @@
         }
     }
 
-    /** This makes a potentially expensive binder call and should be run on a background thread. */
+    /**
+     * This makes a potentially expensive binder call and should be run on a background thread.
+     */
+    @WorkerThread
     public List<StatusBarNotification> getNotificationsForKeys(List<NotificationKeyData> keys) {
-        StatusBarNotification[] notifications = NotificationListener.this
-                .getActiveNotifications(NotificationKeyData.extractKeysOnly(keys)
-                        .toArray(new String[keys.size()]));
-        return notifications == null
-                ? Collections.<StatusBarNotification>emptyList() : Arrays.asList(notifications);
+        StatusBarNotification[] notifications = getActiveNotifications(
+                keys.stream().map(n -> n.notificationKey).toArray(String[]::new));
+        return notifications == null ? Collections.emptyList() : Arrays.asList(notifications);
     }
 
     /**
-     * Filter out notifications that don't have an intent
-     * or are headers for grouped notifications.
-     *
-     * @see #shouldBeFilteredOut(StatusBarNotification)
+     * Returns true for notifications that don't have an intent
+     * or are headers for grouped notifications and should be filtered out.
      */
-    private List<StatusBarNotification> filterNotifications(
-            StatusBarNotification[] notifications) {
-        if (notifications == null) return null;
-        IntSet removedNotifications = new IntSet();
-        for (int i = 0; i < notifications.length; i++) {
-            if (shouldBeFilteredOut(notifications[i])) {
-                removedNotifications.add(i);
-            }
-        }
-        List<StatusBarNotification> filteredNotifications = new ArrayList<>(
-                notifications.length - removedNotifications.size());
-        for (int i = 0; i < notifications.length; i++) {
-            if (!removedNotifications.contains(i)) {
-                filteredNotifications.add(notifications[i]);
-            }
-        }
-        return filteredNotifications;
-    }
-
+    @WorkerThread
     private boolean shouldBeFilteredOut(StatusBarNotification sbn) {
         Notification notification = sbn.getNotification();
-
         updateGroupKeyIfNecessary(sbn);
 
         getCurrentRanking().getRanking(sbn.getKey(), mTempRanking);
@@ -385,16 +331,16 @@
         return (isGroupHeader || missingTitleAndText);
     }
 
+    private static Pair<PackageUserKey, NotificationKeyData> toKeyPair(StatusBarNotification sbn) {
+        return Pair.create(PackageUserKey.fromNotification(sbn),
+                NotificationKeyData.fromNotification(sbn));
+    }
+
     public interface NotificationsChangedListener {
         void onNotificationPosted(PackageUserKey postedPackageUserKey,
-                NotificationKeyData notificationKey, boolean shouldBeFilteredOut);
+                NotificationKeyData notificationKey);
         void onNotificationRemoved(PackageUserKey removedPackageUserKey,
                 NotificationKeyData notificationKey);
         void onNotificationFullRefresh(List<StatusBarNotification> activeNotifications);
     }
-
-    public interface StatusBarNotificationsChangedListener {
-        void onNotificationPosted(StatusBarNotification sbn);
-        void onNotificationRemoved(StatusBarNotification sbn);
-    }
 }
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index 4612b2a..c5aa836 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -20,13 +20,15 @@
 import android.service.notification.StatusBarNotification;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.dot.DotInfo;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.notification.NotificationKeyData;
 import com.android.launcher3.notification.NotificationListener;
-import com.android.launcher3.shortcuts.DeepShortcutManager;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.ShortcutUtil;
@@ -39,13 +41,9 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
 /**
  * Provides data for the popup menu that appears after long-clicking on apps.
  */
@@ -76,28 +74,14 @@
 
     @Override
     public void onNotificationPosted(PackageUserKey postedPackageUserKey,
-            NotificationKeyData notificationKey, boolean shouldBeFilteredOut) {
+            NotificationKeyData notificationKey) {
         DotInfo dotInfo = mPackageUserToDotInfos.get(postedPackageUserKey);
-        boolean dotShouldBeRefreshed;
         if (dotInfo == null) {
-            if (!shouldBeFilteredOut) {
-                DotInfo newDotInfo = new DotInfo();
-                newDotInfo.addOrUpdateNotificationKey(notificationKey);
-                mPackageUserToDotInfos.put(postedPackageUserKey, newDotInfo);
-                dotShouldBeRefreshed = true;
-            } else {
-                dotShouldBeRefreshed = false;
-            }
-        } else {
-            dotShouldBeRefreshed = shouldBeFilteredOut
-                    ? dotInfo.removeNotificationKey(notificationKey)
-                    : dotInfo.addOrUpdateNotificationKey(notificationKey);
-            if (dotInfo.getNotificationKeys().size() == 0) {
-                mPackageUserToDotInfos.remove(postedPackageUserKey);
-            }
+            dotInfo = new DotInfo();
+            mPackageUserToDotInfos.put(postedPackageUserKey, dotInfo);
         }
-        if (dotShouldBeRefreshed) {
-            updateNotificationDots(t -> postedPackageUserKey.equals(t));
+        if (dotInfo.addOrUpdateNotificationKey(notificationKey)) {
+            updateNotificationDots(postedPackageUserKey::equals);
         }
     }
 
@@ -109,7 +93,7 @@
             if (oldDotInfo.getNotificationKeys().size() == 0) {
                 mPackageUserToDotInfos.remove(removedPackageUserKey);
             }
-            updateNotificationDots(t -> removedPackageUserKey.equals(t));
+            updateNotificationDots(removedPackageUserKey::equals);
             trimNotifications(mPackageUserToDotInfos);
         }
     }
@@ -195,14 +179,6 @@
                 : getNotificationsForItem(info, dotInfo.getNotificationKeys());
     }
 
-    /** This makes a potentially expensive binder call and should be run on a background thread. */
-    public @NonNull List<StatusBarNotification> getStatusBarNotificationsForKeys(
-            List<NotificationKeyData> notificationKeys) {
-        NotificationListener notificationListener = NotificationListener.getInstanceIfConnected();
-        return notificationListener == null ? Collections.EMPTY_LIST
-                : notificationListener.getNotificationsForKeys(notificationKeys);
-    }
-
     public void cancelNotification(String notificationKey) {
         NotificationListener notificationListener = NotificationListener.getInstanceIfConnected();
         if (notificationListener == null) {
diff --git a/src/com/android/launcher3/popup/PopupPopulator.java b/src/com/android/launcher3/popup/PopupPopulator.java
index dbfe988..9618927 100644
--- a/src/com/android/launcher3/popup/PopupPopulator.java
+++ b/src/com/android/launcher3/popup/PopupPopulator.java
@@ -20,7 +20,9 @@
 import android.content.pm.ShortcutInfo;
 import android.os.Handler;
 import android.os.UserHandle;
-import android.service.notification.StatusBarNotification;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.Launcher;
@@ -28,6 +30,7 @@
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.notification.NotificationInfo;
 import com.android.launcher3.notification.NotificationKeyData;
+import com.android.launcher3.notification.NotificationListener;
 import com.android.launcher3.shortcuts.DeepShortcutManager;
 import com.android.launcher3.shortcuts.DeepShortcutView;
 import com.android.launcher3.util.PackageUserKey;
@@ -37,9 +40,7 @@
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
-
-import androidx.annotation.Nullable;
-import androidx.annotation.VisibleForTesting;
+import java.util.stream.Collectors;
 
 /**
  * Contains logic relevant to populating a {@link PopupContainerWithArrow}. In particular,
@@ -130,12 +131,15 @@
         final UserHandle user = originalInfo.user;
         return () -> {
             if (!notificationKeys.isEmpty()) {
-                List<StatusBarNotification> notifications = launcher.getPopupDataProvider()
-                        .getStatusBarNotificationsForKeys(notificationKeys);
-                List<NotificationInfo> infos = new ArrayList<>(notifications.size());
-                for (int i = 0; i < notifications.size(); i++) {
-                    StatusBarNotification notification = notifications.get(i);
-                    infos.add(new NotificationInfo(launcher, notification));
+                NotificationListener notificationListener =
+                        NotificationListener.getInstanceIfConnected();
+                final List<NotificationInfo> infos;
+                if (notificationListener == null) {
+                    infos = Collections.emptyList();
+                } else {
+                    infos = notificationListener.getNotificationsForKeys(notificationKeys).stream()
+                            .map(sbn -> new NotificationInfo(launcher, sbn))
+                            .collect(Collectors.toList());
                 }
                 uiHandler.post(() -> container.applyNotificationInfos(infos));
             }
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index 140a06a..64df384 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -24,9 +24,10 @@
 import android.graphics.Color;
 import android.os.Bundle;
 import android.os.Debug;
-import android.util.Log;
 import android.view.View;
 
+import androidx.annotation.Keep;
+
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.Launcher;
@@ -186,6 +187,22 @@
         Runtime.getRuntime().runFinalization();
 
         final CountDownLatch fence = new CountDownLatch(1);
+        createFinalizationObserver(fence);
+        try {
+            do {
+                Runtime.getRuntime().gc();
+                Runtime.getRuntime().runFinalization();
+            } while (!fence.await(100, TimeUnit.MILLISECONDS));
+        } catch (InterruptedException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
+
+    // Create the observer in the scope of a method to minimize the chance that
+    // it remains live in a DEX/machine register at the point of the fence guard.
+    // This must be kept to avoid R8 inlining it.
+    @Keep
+    private static void createFinalizationObserver(CountDownLatch fence) {
         new Object() {
             @Override
             protected void finalize() throws Throwable {
@@ -196,13 +213,5 @@
                 }
             }
         };
-        try {
-            do {
-                Runtime.getRuntime().gc();
-                Runtime.getRuntime().runFinalization();
-            } while (!fence.await(100, TimeUnit.MILLISECONDS));
-        } catch (InterruptedException ex) {
-            throw new RuntimeException(ex);
-        }
     }
 }
diff --git a/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java b/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java
index 08fa098..858cb38 100644
--- a/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java
+++ b/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java
@@ -21,6 +21,7 @@
 import android.os.Build;
 import android.util.Log;
 
+import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.UiDevice;
 
 import org.junit.rules.TestRule;
@@ -86,6 +87,19 @@
     }
 
     private static int getRunFlavor() {
+        final String flavorOverride = InstrumentationRegistry.getArguments().getString("flavor");
+
+        if (flavorOverride != null) {
+            Log.d(TAG, "Flavor override: " + flavorOverride);
+            try {
+                return (int) TestStabilityRule.class.getField(flavorOverride).get(null);
+            } catch (NoSuchFieldException e) {
+                throw new AssertionError("Unrecognized run flavor override: " + flavorOverride);
+            } catch (IllegalAccessException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
         final String launcherVersion;
         try {
             launcherVersion = getInstrumentation().