diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatEduController.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatEduController.java
index 63e7390..680012c 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatEduController.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatEduController.java
@@ -30,6 +30,7 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.data.FolderInfo;
@@ -298,7 +299,7 @@
             Log.e(TAG, "Unable to find suitable view for ArrowTip");
             return false;
         }
-        Rect bounds = mLauncher.getViewBounds(tipTargetView);
+        Rect bounds = Utilities.getViewBounds(tipTargetView);
         new ArrowTipView(mLauncher).show(message, Gravity.END, bounds.centerX(), bounds.top);
         return true;
     }
diff --git a/quickstep/src/com/android/launcher3/model/WellbeingModel.java b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
index 154b78b..e489cb3 100644
--- a/quickstep/src/com/android/launcher3/model/WellbeingModel.java
+++ b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
@@ -377,7 +377,7 @@
     /**
      * Shortcut factory for generating wellbeing action
      */
-    public static final SystemShortcut.Factory SHORTCUT_FACTORY =
+    public static final SystemShortcut.Factory<BaseDraggingActivity> SHORTCUT_FACTORY =
             (activity, info) -> (info.getTargetComponent() == null) ? null : INSTANCE.get(activity)
                     .getShortcutForApp(
                             info.getTargetComponent().getPackageName(), info.user.getIdentifier(),
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index f1d7d41..370496a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -70,10 +70,10 @@
 import com.android.quickstep.SysUINavigationMode;
 import com.android.quickstep.SysUINavigationMode.Mode;
 import com.android.quickstep.SystemUiProxy;
-import com.android.quickstep.util.ScopedUnfoldTransitionProgressProvider;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.WindowManagerWrapper;
+import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider;
 
 /**
  * The {@link ActivityContext} with which we inflate Taskbar-related Views. This allows UI elements
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 63d07f3..5bedf87 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -47,7 +47,7 @@
 import com.android.quickstep.SysUINavigationMode.Mode;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TouchInteractionService;
-import com.android.quickstep.util.ScopedUnfoldTransitionProgressProvider;
+import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider;
 
 /**
  * Class to manage taskbar lifecycle
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUnfoldAnimationController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUnfoldAnimationController.java
index 978bd47..c785186 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUnfoldAnimationController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUnfoldAnimationController.java
@@ -19,9 +19,9 @@
 import android.view.WindowManager;
 
 import com.android.quickstep.util.LauncherViewsMoveFromCenterTranslationApplier;
-import com.android.quickstep.util.ScopedUnfoldTransitionProgressProvider;
 import com.android.systemui.shared.animation.UnfoldMoveFromCenterAnimator;
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener;
+import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider;
 
 /**
  * Controls animation of taskbar icons when unfolding foldable devices
diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index dcc7ccc..8c4ba97 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -128,7 +128,7 @@
         }
     }
 
-    class MultiWindowSystemShortcut extends SystemShortcut {
+    class MultiWindowSystemShortcut extends SystemShortcut<BaseDraggingActivity> {
 
         private Handler mHandler;
 
@@ -305,7 +305,7 @@
         return new PinSystemShortcut(activity, taskContainer);
     };
 
-    class PinSystemShortcut extends SystemShortcut {
+    class PinSystemShortcut extends SystemShortcut<BaseDraggingActivity> {
 
         private static final String TAG = "PinSystemShortcut";
 
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 4844f6b..ddcf34c 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -145,6 +145,8 @@
 
     private int mBackGestureNotificationCounter = -1;
 
+    private final TISBinder mTISBinder = new TISBinder();
+
     /**
      * Local IOverviewProxy implementation with some methods for local components
      */
@@ -516,7 +518,7 @@
     @Override
     public IBinder onBind(Intent intent) {
         Log.d(TAG, "Touch service connected: user=" + getUserId());
-        return new TISBinder();
+        return mTISBinder;
     }
 
     private void onInputEvent(InputEvent ev) {
diff --git a/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java b/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java
index 47d3580..b39412b 100644
--- a/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java
+++ b/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java
@@ -26,6 +26,7 @@
 import com.android.launcher3.util.HorizontalInsettableView;
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener;
+import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider;
 
 /**
  * Controls animations that are happening during unfolding foldable devices
diff --git a/quickstep/src/com/android/quickstep/util/ScopedUnfoldTransitionProgressProvider.java b/quickstep/src/com/android/quickstep/util/ScopedUnfoldTransitionProgressProvider.java
deleted file mode 100644
index 2ef311f..0000000
--- a/quickstep/src/com/android/quickstep/util/ScopedUnfoldTransitionProgressProvider.java
+++ /dev/null
@@ -1,140 +0,0 @@
-/*
- * Copyright (C) 2021 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.quickstep.util;
-
-import android.annotation.NonNull;
-import android.annotation.Nullable;
-
-import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
-import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * Manages progress listeners that can have smaller lifespan than the unfold animation.
- * Allows to limit getting transition updates to only when
- * {@link ScopedUnfoldTransitionProgressProvider#setReadyToHandleTransition} is called
- * with readyToHandleTransition = true
- *
- * If the transition has already started by the moment when the clients are ready to play
- * the transition then it will report transition started callback and current animation progress.
- */
-public final class ScopedUnfoldTransitionProgressProvider implements
-        UnfoldTransitionProgressProvider, TransitionProgressListener {
-
-    private static final float PROGRESS_UNSET = -1f;
-
-    @Nullable
-    private UnfoldTransitionProgressProvider mSource;
-
-    private final List<TransitionProgressListener> mListeners = new ArrayList<>();
-
-    private boolean mIsReadyToHandleTransition;
-    private boolean mIsTransitionRunning;
-    private float mLastTransitionProgress = PROGRESS_UNSET;
-
-    public ScopedUnfoldTransitionProgressProvider() {
-        this(null);
-    }
-
-    public ScopedUnfoldTransitionProgressProvider(@Nullable UnfoldTransitionProgressProvider
-                                                          source) {
-        setSourceProvider(source);
-    }
-
-    /**
-     * Sets the source for the unfold transition progress updates,
-     * it replaces current provider if it is already set
-     * @param provider transition provider that emits transition progress updates
-     */
-    public void setSourceProvider(@Nullable UnfoldTransitionProgressProvider provider) {
-        if (mSource != null) {
-            mSource.removeCallback(this);
-        }
-
-        if (provider != null) {
-            mSource = provider;
-            mSource.addCallback(this);
-        }
-    }
-
-    /**
-     * Allows to notify this provide whether the listeners can play the transition or not.
-     * Call this method with readyToHandleTransition = true when all listeners
-     * are ready to consume the transition progress events.
-     * Call it with readyToHandleTransition = false when listeners can't process the events.
-     */
-    public void setReadyToHandleTransition(boolean isReadyToHandleTransition) {
-        if (mIsTransitionRunning) {
-            if (mIsReadyToHandleTransition) {
-                mListeners.forEach(TransitionProgressListener::onTransitionStarted);
-
-                if (mLastTransitionProgress != PROGRESS_UNSET) {
-                    mListeners.forEach(listener ->
-                            listener.onTransitionProgress(mLastTransitionProgress));
-                }
-            } else {
-                mIsTransitionRunning = false;
-                mListeners.forEach(TransitionProgressListener::onTransitionFinished);
-            }
-        }
-
-        mIsReadyToHandleTransition = isReadyToHandleTransition;
-    }
-
-    @Override
-    public void addCallback(@NonNull TransitionProgressListener listener) {
-        mListeners.add(listener);
-    }
-
-    @Override
-    public void removeCallback(@NonNull TransitionProgressListener listener) {
-        mListeners.remove(listener);
-    }
-
-    @Override
-    public void destroy() {
-        mSource.removeCallback(this);
-    }
-
-    @Override
-    public void onTransitionStarted() {
-        this.mIsTransitionRunning = true;
-        if (mIsReadyToHandleTransition) {
-            mListeners.forEach(TransitionProgressListener::onTransitionStarted);
-        }
-    }
-
-    @Override
-    public void onTransitionProgress(float progress) {
-        if (mIsReadyToHandleTransition) {
-            mListeners.forEach(listener -> listener.onTransitionProgress(progress));
-        }
-
-        mLastTransitionProgress = progress;
-    }
-
-    @Override
-    public void onTransitionFinished() {
-        if (mIsReadyToHandleTransition) {
-            mListeners.forEach(TransitionProgressListener::onTransitionFinished);
-        }
-
-        mIsTransitionRunning = false;
-        mLastTransitionProgress = PROGRESS_UNSET;
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 292e9d7..f7a9562 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -343,6 +343,7 @@
                         }
                     });
                     view.setTaskViewsResistanceTranslation(view.mTaskViewsSecondaryTranslation);
+                    view.updateTaskViewsSnapshotRadius();
                     view.updatePageOffsets();
                 }
 
@@ -3741,6 +3742,12 @@
                         .recentsViewSecondaryTranslation.value = translation);
     }
 
+    private void updateTaskViewsSnapshotRadius() {
+        for (int i = 0; i < getTaskViewCount(); i++) {
+            getTaskViewAt(i).updateSnapshotRadius();
+        }
+    }
+
     protected void setTaskViewsPrimarySplitTranslation(float translation) {
         mTaskViewsPrimarySplitTranslation = translation;
         for (int i = 0; i < getTaskViewCount(); i++) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 31a73e9..8d77e44 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -1341,7 +1341,7 @@
         invalidateOutline();
     }
 
-    private void updateSnapshotRadius() {
+    void updateSnapshotRadius() {
         updateCurrentFullscreenParams(mSnapshotView.getPreviewPositionHelper());
         mSnapshotView.setFullscreenParams(mCurrentFullscreenParams);
     }
diff --git a/src/com/android/launcher3/AppWidgetResizeFrame.java b/src/com/android/launcher3/AppWidgetResizeFrame.java
index ebfd281..1bfd7b5 100644
--- a/src/com/android/launcher3/AppWidgetResizeFrame.java
+++ b/src/com/android/launcher3/AppWidgetResizeFrame.java
@@ -215,7 +215,7 @@
 
         dl.addView(frame);
         frame.mIsOpen = true;
-        frame.snapToWidget(false);
+        frame.post(() -> frame.snapToWidget(false));
     }
 
     private void setupForWidget(LauncherAppWidgetHostView widgetView, CellLayout cellLayout,
diff --git a/src/com/android/launcher3/BaseDraggingActivity.java b/src/com/android/launcher3/BaseDraggingActivity.java
index 7954011..dd56ca3 100644
--- a/src/com/android/launcher3/BaseDraggingActivity.java
+++ b/src/com/android/launcher3/BaseDraggingActivity.java
@@ -41,7 +41,6 @@
 import android.view.ActionMode;
 import android.view.Display;
 import android.view.View;
-import android.view.View.OnClickListener;
 import android.view.WindowInsets.Type;
 import android.view.WindowMetrics;
 import android.widget.Toast;
@@ -166,12 +165,6 @@
         // no-op
     }
 
-    public Rect getViewBounds(View v) {
-        int[] pos = new int[2];
-        v.getLocationOnScreen(pos);
-        return new Rect(pos[0], pos[1], pos[0] + v.getWidth(), pos[1] + v.getHeight());
-    }
-
     @NonNull
     public ActivityOptionsWrapper getActivityLaunchOptions(View v, @Nullable ItemInfo item) {
         int left = 0, top = 0;
@@ -206,7 +199,7 @@
         // Prepare intent
         intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
         if (v != null) {
-            intent.setSourceBounds(getViewBounds(v));
+            intent.setSourceBounds(Utilities.getViewBounds(v));
         }
         try {
             boolean isShortcut = (item instanceof WorkspaceItemInfo)
@@ -316,7 +309,8 @@
         }
     }
 
-    public OnClickListener getItemOnClickListener() {
+    @Override
+    public View.OnClickListener getItemOnClickListener() {
         return ItemClickHandler.INSTANCE;
     }
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index f429d76..8d92bf2 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1391,22 +1391,7 @@
             final LauncherAppWidgetHostView launcherHostView = (LauncherAppWidgetHostView) hostView;
             CellLayout cellLayout = getCellLayout(launcherInfo.container, launcherInfo.screenId);
             if (mStateManager.getState() == NORMAL) {
-                // Show resize frame once the widget layout is drawn.
-                View.OnLayoutChangeListener onLayoutChangeListener =
-                        new View.OnLayoutChangeListener() {
-                            @Override
-                            public void onLayoutChange(View view, int left, int top, int right,
-                                    int bottom, int oldLeft, int oldTop, int oldRight,
-                                    int oldBottom) {
-                                AppWidgetResizeFrame.showForWidget(launcherHostView, cellLayout);
-                                launcherHostView.removeOnLayoutChangeListener(this);
-                            }
-                        };
-                launcherHostView.addOnLayoutChangeListener(onLayoutChangeListener);
-                // There is a small chance that the layout was already drawn before the layout
-                // change listener was registered, which means that the resize frame wouldn't be
-                // shown. Directly call requestLayout to force a layout change.
-                launcherHostView.requestLayout();
+                AppWidgetResizeFrame.showForWidget(launcherHostView, cellLayout);
             } else {
                 mStateManager.addStateListener(new StateManager.StateListener<LauncherState>() {
                     @Override
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 36faeee..7a38fe7 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -69,6 +69,7 @@
 import android.view.animation.Interpolator;
 import android.widget.LinearLayout;
 
+import androidx.annotation.NonNull;
 import androidx.core.graphics.ColorUtils;
 import androidx.core.os.BuildCompat;
 
@@ -846,6 +847,12 @@
         view.setLayoutParams(lp);
     }
 
+    public static Rect getViewBounds(@NonNull View v) {
+        int[] pos = new int[2];
+        v.getLocationOnScreen(pos);
+        return new Rect(pos[0], pos[1], pos[0] + v.getWidth(), pos[1] + v.getHeight());
+    }
+
     private static class FixedSizeEmptyDrawable extends ColorDrawable {
 
         private final int mSize;
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 8e76d82..8095280 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -1676,7 +1676,7 @@
         }
 
         if (child instanceof BubbleTextView && !dragOptions.isAccessibleDrag) {
-            PopupContainerWithArrow popupContainer = PopupContainerWithArrow
+            PopupContainerWithArrow<Launcher> popupContainer = PopupContainerWithArrow
                     .showForIcon((BubbleTextView) child);
             if (popupContainer != null) {
                 dragOptions.preDragCondition = popupContainer.createPreDragCondition();
diff --git a/src/com/android/launcher3/popup/PopupContainerWithArrow.java b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
index d6e927b..b963950 100644
--- a/src/com/android/launcher3/popup/PopupContainerWithArrow.java
+++ b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
@@ -66,6 +66,7 @@
 import com.android.launcher3.touch.ItemLongClickListener;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.ShortcutUtil;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 
 import java.util.ArrayList;
@@ -81,7 +82,7 @@
  *
  * @param <T> The activity on with the popup shows
  */
-public class PopupContainerWithArrow<T extends BaseDraggingActivity>
+public class PopupContainerWithArrow<T extends Context & ActivityContext>
         extends ArrowPopup<T> implements DragSource, DragController.DragListener {
 
     private final List<DeepShortcutView> mShortcuts = new ArrayList<>();
@@ -190,10 +191,10 @@
     }
 
     /**
-     * Shows the notifications and deep shortcuts associated with {@param icon}.
+     * Shows the notifications and deep shortcuts associated with a Launcher {@param icon}.
      * @return the container if shown or null.
      */
-    public static PopupContainerWithArrow showForIcon(BubbleTextView icon) {
+    public static PopupContainerWithArrow<Launcher> showForIcon(BubbleTextView icon) {
         Launcher launcher = Launcher.getLauncher(icon.getContext());
         if (getOpen(launcher) != null) {
             // There is already an items container open, so don't open this one.
@@ -205,7 +206,7 @@
             return null;
         }
 
-        final PopupContainerWithArrow container =
+        final PopupContainerWithArrow<Launcher> container =
                 (PopupContainerWithArrow) launcher.getLayoutInflater().inflate(
                         R.layout.popup_container, launcher.getDragLayer(), false);
         container.configureForLauncher(launcher);
@@ -489,8 +490,8 @@
     /**
      * Returns a PopupContainerWithArrow which is already open or null
      */
-    public static PopupContainerWithArrow getOpen(BaseDraggingActivity launcher) {
-        return getOpenView(launcher, TYPE_ACTION_POPUP);
+    public static <T extends Context & ActivityContext> PopupContainerWithArrow getOpen(T context) {
+        return getOpenView(context, TYPE_ACTION_POPUP);
     }
 
     /**
diff --git a/src/com/android/launcher3/popup/PopupPopulator.java b/src/com/android/launcher3/popup/PopupPopulator.java
index 5ed6f2e..1dce1f2 100644
--- a/src/com/android/launcher3/popup/PopupPopulator.java
+++ b/src/com/android/launcher3/popup/PopupPopulator.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_SHORTCUTS;
 
 import android.content.ComponentName;
+import android.content.Context;
 import android.content.pm.ShortcutInfo;
 import android.os.Handler;
 import android.os.UserHandle;
@@ -26,7 +27,6 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
-import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.ItemInfo;
@@ -36,6 +36,7 @@
 import com.android.launcher3.notification.NotificationListener;
 import com.android.launcher3.shortcuts.DeepShortcutView;
 import com.android.launcher3.shortcuts.ShortcutRequest;
+import com.android.launcher3.views.ActivityContext;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -128,7 +129,8 @@
     /**
      * Returns a runnable to update the provided shortcuts and notifications
      */
-    public static Runnable createUpdateRunnable(final BaseDraggingActivity launcher,
+    public static <T extends Context & ActivityContext> Runnable createUpdateRunnable(
+            final T context,
             final ItemInfo originalInfo,
             final Handler uiHandler, final PopupContainerWithArrow container,
             final List<DeepShortcutView> shortcutViews,
@@ -144,22 +146,22 @@
                     infos = Collections.emptyList();
                 } else {
                     infos = notificationListener.getNotificationsForKeys(notificationKeys).stream()
-                            .map(sbn -> new NotificationInfo(launcher, sbn, originalInfo))
+                            .map(sbn -> new NotificationInfo(context, sbn, originalInfo))
                             .collect(Collectors.toList());
                 }
                 uiHandler.post(() -> container.applyNotificationInfos(infos));
             }
 
-            List<ShortcutInfo> shortcuts = new ShortcutRequest(launcher, user)
+            List<ShortcutInfo> shortcuts = new ShortcutRequest(context, user)
                     .withContainer(activity)
                     .query(ShortcutRequest.PUBLISHED);
             String shortcutIdToDeDupe = notificationKeys.isEmpty() ? null
                     : notificationKeys.get(0).shortcutId;
             shortcuts = PopupPopulator.sortAndFilterShortcuts(shortcuts, shortcutIdToDeDupe);
-            IconCache cache = LauncherAppState.getInstance(launcher).getIconCache();
+            IconCache cache = LauncherAppState.getInstance(context).getIconCache();
             for (int i = 0; i < shortcuts.size() && i < shortcutViews.size(); i++) {
                 final ShortcutInfo shortcut = shortcuts.get(i);
-                final WorkspaceItemInfo si = new WorkspaceItemInfo(shortcut, launcher);
+                final WorkspaceItemInfo si = new WorkspaceItemInfo(shortcut, context);
                 cache.getUnbadgedShortcutIcon(si, shortcut);
                 si.rank = i;
                 si.container = CONTAINER_SHORTCUTS;
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index e5424cf..826c79b 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -18,12 +18,14 @@
 import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.InstantAppResolver;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.WidgetsBottomSheet;
 
 import java.util.List;
@@ -35,7 +37,7 @@
  * Example system shortcuts, defined as inner classes, include Widgets and AppInfo.
  * @param <T>
  */
-public abstract class SystemShortcut<T extends BaseDraggingActivity> extends ItemInfo
+public abstract class SystemShortcut<T extends Context & ActivityContext> extends ItemInfo
         implements View.OnClickListener {
 
     private final int mIconResId;
@@ -100,7 +102,7 @@
         return mAccessibilityActionId == action;
     }
 
-    public interface Factory<T extends BaseDraggingActivity> {
+    public interface Factory<T extends Context & ActivityContext> {
 
         @Nullable SystemShortcut<T> getShortcut(T activity, ItemInfo itemInfo);
     }
@@ -135,9 +137,9 @@
 
     public static final Factory<BaseDraggingActivity> APP_INFO = AppInfo::new;
 
-    public static class AppInfo extends SystemShortcut {
+    public static class AppInfo<T extends Context & ActivityContext> extends SystemShortcut<T> {
 
-        public AppInfo(BaseDraggingActivity target, ItemInfo itemInfo) {
+        public AppInfo(T target, ItemInfo itemInfo) {
             super(R.drawable.ic_info_no_shadow, R.string.app_info_drop_target_label, target,
                     itemInfo);
         }
@@ -145,7 +147,7 @@
         @Override
         public void onClick(View view) {
             dismissTaskMenuView(mTarget);
-            Rect sourceBounds = mTarget.getViewBounds(view);
+            Rect sourceBounds = Utilities.getViewBounds(view);
             new PackageManagerHelper(mTarget).startDetailsActivityForInfo(
                     mItemInfo, sourceBounds, ActivityOptions.makeBasic().toBundle());
             mTarget.getStatsLogManager().logger().withItemInfo(mItemInfo)
@@ -170,7 +172,7 @@
         return new Install(activity, itemInfo);
     };
 
-    public static class Install extends SystemShortcut {
+    public static class Install extends SystemShortcut<BaseDraggingActivity> {
 
         public Install(BaseDraggingActivity target, ItemInfo itemInfo) {
             super(R.drawable.ic_install_no_shadow, R.string.install_drop_target_label,
@@ -186,7 +188,7 @@
         }
     }
 
-    public static void dismissTaskMenuView(BaseDraggingActivity activity) {
+    public static <T extends Context & ActivityContext> void dismissTaskMenuView(T activity) {
         AbstractFloatingView.closeOpenViews(activity, true,
             AbstractFloatingView.TYPE_ALL & ~AbstractFloatingView.TYPE_REBIND_SAFE);
     }
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index e07d71e..a2e4ad6 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -19,6 +19,7 @@
 import android.content.ContextWrapper;
 import android.graphics.Rect;
 import android.view.LayoutInflater;
+import android.view.View;
 import android.view.View.AccessibilityDelegate;
 
 import androidx.annotation.Nullable;
@@ -159,4 +160,10 @@
             return null;
         }
     }
+
+    default View.OnClickListener getItemOnClickListener() {
+        return v -> {
+            // No op.
+        };
+    }
 }
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 5ea5d65..2c9785c 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -15,28 +15,24 @@
  */
 package com.android.launcher3.ui.widget;
 
-import static androidx.test.InstrumentationRegistry.getInstrumentation;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
 
 import android.appwidget.AppWidgetManager;
 import android.content.Intent;
-import android.view.View;
 
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.tapl.Widget;
+import com.android.launcher3.tapl.WidgetResizeFrame;
 import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.testcomponent.WidgetConfigActivity;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
-import com.android.launcher3.util.LauncherBindableItemsContainer.ItemOperator;
-import com.android.launcher3.util.Wait;
-import com.android.launcher3.util.Wait.Condition;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
@@ -92,48 +88,26 @@
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        widgets.getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
-                .dragToWorkspace(true, false);
+        WidgetResizeFrame resizeFrame =
+                widgets.getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                        .dragConfigWidgetToWorkspace(acceptConfig);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
 
         // Verify that the widget id is valid and bound
         assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
 
-        setResult(acceptConfig);
         if (acceptConfig) {
-            // TODO(b/192655785) Assert widget resize frame is shown and then dismiss it.
-            Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
-            assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
+            assertNotNull("Widget resize frame not shown after widget added", resizeFrame);
+            resizeFrame.dismiss();
+
+            final Widget widget =
+                    mLauncher.getWorkspace().tryGetWidget(mWidgetInfo.label, DEFAULT_UI_TIMEOUT);
+            assertNotNull("Widget not found on the workspace", widget);
         } else {
-            // Verify that the widget id is deleted.
-            Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
-                    DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
-        }
-    }
-
-    private void setResult(boolean success) {
-        getInstrumentation().getTargetContext().sendBroadcast(
-                WidgetConfigActivity.getCommandIntent(WidgetConfigActivity.class,
-                        success ? "clickOK" : "clickCancel"));
-    }
-
-    /**
-     * Condition for searching widget id
-     */
-    private class WidgetSearchCondition implements Condition, ItemOperator {
-
-        @Override
-        public boolean isTrue() throws Throwable {
-            return mMainThreadExecutor.submit(mActivityMonitor.itemExists(this)).get();
-        }
-
-        @Override
-        public boolean evaluate(ItemInfo info, View view) {
-            return info instanceof LauncherAppWidgetInfo &&
-                    ((LauncherAppWidgetInfo) info).providerName.getClassName().equals(
-                            mWidgetInfo.provider.getClassName()) &&
-                    ((LauncherAppWidgetInfo) info).appWidgetId == mWidgetId;
+            final Widget widget =
+                    mLauncher.getWorkspace().tryGetWidget(mWidgetInfo.label, DEFAULT_UI_TIMEOUT);
+            assertNull("Widget unexpectedly found on the workspace", widget);
         }
     }
 
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index dad4f2b..194ee4f 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -25,6 +25,7 @@
 
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.tapl.Widget;
+import com.android.launcher3.tapl.WidgetResizeFrame;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
 import com.android.launcher3.util.rule.ShellCommandRule;
@@ -53,19 +54,20 @@
         final LauncherAppWidgetProviderInfo widgetInfo =
                 TestViewHelpers.findWidgetProvider(this, false /* hasConfigureScreen */);
 
-        mLauncher.
+        WidgetResizeFrame resizeFrame = mLauncher.
                 getWorkspace().
                 openAllWidgets().
                 getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace(false, false);
-        // Dismiss widget resize frame.
-        mDevice.pressHome();
+                dragWidgetToWorkspace();
 
         assertTrue(mActivityMonitor.itemExists(
                 (info, view) -> info instanceof LauncherAppWidgetInfo &&
                         ((LauncherAppWidgetInfo) info).providerName.getClassName().equals(
                                 widgetInfo.provider.getClassName())).call());
 
+        assertNotNull("Widget resize frame not shown after widget add", resizeFrame);
+        resizeFrame.dismiss();
+
         final Widget widget = mLauncher.getWorkspace().tryGetWidget(widgetInfo.label,
                 DEFAULT_UI_TIMEOUT);
         assertNotNull("Widget not found on the workspace", widget);
diff --git a/tests/tapl/com/android/launcher3/tapl/Widget.java b/tests/tapl/com/android/launcher3/tapl/Widget.java
index 3520318..f569ef4 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widget.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widget.java
@@ -16,7 +16,12 @@
 
 package com.android.launcher3.tapl;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiObject2;
+import androidx.test.uiautomator.Until;
 
 import com.android.launcher3.testing.TestProtocol;
 
@@ -51,4 +56,55 @@
     protected String launchableType() {
         return "widget";
     }
+
+    /**
+     * Drags a non-configurable widget from the widgets container to the workspace and returns the
+     * resize frame that is shown after the widget is added.
+     */
+    @NonNull
+    public WidgetResizeFrame dragWidgetToWorkspace() {
+        return dragWidgetToWorkspace(/* configurable= */ false, /* acceptsConfig= */ false);
+    }
+
+    /**
+     * Drags a configurable widget from the widgets container to the workspace, either accepts or
+     * cancels the configuration based on {@code acceptsConfig}, and returns the resize frame that
+     * is shown if the widget is added.
+     */
+    @Nullable
+    public WidgetResizeFrame dragConfigWidgetToWorkspace(boolean acceptsConfig) {
+        return dragWidgetToWorkspace(/* configurable= */ true, acceptsConfig);
+    }
+
+    /**
+     * Drags a widget from the widgets container to the workspace and returns the resize frame that
+     * is shown after the widget is added.
+     *
+     * <p> If {@code configurable} is true, then either accepts or cancels the configuration based
+     * on {@code acceptsConfig}.
+     */
+    @Nullable
+    private WidgetResizeFrame dragWidgetToWorkspace(
+            boolean configurable, boolean acceptsConfig) {
+        dragToWorkspace(/* startsActivity= */ configurable, /* isWidgetShortcut= */ false);
+
+        if (configurable) {
+            // Configure the widget.
+            BySelector selector = By.text(acceptsConfig ? "OK" : "Cancel");
+            mLauncher.getDevice()
+                    .wait(Until.findObject(selector), LauncherInstrumentation.WAIT_TIME_MS)
+                    .click();
+
+            // If the widget configuration was cancelled, then the widget wasn't added to the home
+            // screen. In that case, we cannot return a resize frame.
+            if (!acceptsConfig) {
+                return null;
+            }
+        }
+
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to get widget resize frame")) {
+            return new WidgetResizeFrame(mLauncher);
+        }
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java b/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java
new file mode 100644
index 0000000..8f51d04
--- /dev/null
+++ b/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2021 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.launcher3.tapl;
+
+/** The resize frame that is shown for a widget on the workspace. */
+public class WidgetResizeFrame {
+
+    private final LauncherInstrumentation mLauncher;
+
+    WidgetResizeFrame(LauncherInstrumentation launcher) {
+        mLauncher = launcher;
+        launcher.waitForLauncherObject("widget_resize_frame");
+    }
+
+    /** Dismisses the resize frame. */
+    public void dismiss() {
+        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
+             LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                     "want to dismiss widget resize frame")) {
+            // Dismiss the resize frame by pressing the home button.
+            mLauncher.getDevice().pressHome();
+        }
+    }
+}