diff --git a/quickstep/AndroidManifest-launcher.xml b/quickstep/AndroidManifest-launcher.xml
index c6e2d8c..80d8154 100644
--- a/quickstep/AndroidManifest-launcher.xml
+++ b/quickstep/AndroidManifest-launcher.xml
@@ -48,7 +48,7 @@
             android:stateNotNeeded="true"
             android:windowSoftInputMode="adjustPan"
             android:screenOrientation="unspecified"
-            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize"
+            android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize|uiMode"
             android:resizeableActivity="true"
             android:resumeWhilePausing="true"
             android:taskAffinity=""
diff --git a/quickstep/AndroidManifest.xml b/quickstep/AndroidManifest.xml
index bf198b6..4abf6e1 100644
--- a/quickstep/AndroidManifest.xml
+++ b/quickstep/AndroidManifest.xml
@@ -80,7 +80,7 @@
              android:stateNotNeeded="true"
              android:theme="@style/LauncherTheme"
              android:screenOrientation="behind"
-             android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize"
+             android:configChanges="keyboard|keyboardHidden|mcc|mnc|navigation|orientation|screenSize|screenLayout|smallestScreenSize|uiMode"
              android:resizeableActivity="true"
              android:resumeWhilePausing="true"
              android:enableOnBackInvokedCallback="false"
diff --git a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
index 23cb8e9..943c08c 100644
--- a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
+++ b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
@@ -194,6 +194,11 @@
                 return false;
             }
 
+            View dragView = widgetCell.getDragAndDropView();
+            if (dragView == null) {
+                return false;
+            }
+
             ClipData clipData = new ClipData(
                     new ClipDescription(
                             /* label= */ "", // not displayed anywhere; so, set to empty.
@@ -209,9 +214,9 @@
                     .putExtra(EXTRA_IS_PENDING_WIDGET_DRAG, true));
 
             // DRAG_FLAG_GLOBAL permits dragging data beyond app window.
-            return view.startDragAndDrop(
+            return dragView.startDragAndDrop(
                     clipData,
-                    new View.DragShadowBuilder(view),
+                    new View.DragShadowBuilder(dragView),
                     /* myLocalState= */ null,
                     View.DRAG_FLAG_GLOBAL
             );
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 39f2c00..64bb05e 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 
 import android.app.prediction.AppTarget;
+import android.content.ComponentName;
 import android.content.Context;
 import android.text.TextUtils;
 
@@ -29,7 +30,6 @@
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.picker.WidgetRecommendationCategoryProvider;
 
@@ -64,36 +64,24 @@
                 widget -> new ComponentKey(widget.providerName, widget.user)).collect(
                 Collectors.toSet());
         Predicate<WidgetItem> notOnWorkspace = w -> !widgetsInWorkspace.contains(w);
-        Map<PackageUserKey, List<WidgetItem>> allWidgets =
-                dataModel.widgetsModel.getAllWidgetsWithoutShortcuts();
+        Map<ComponentKey, WidgetItem> allWidgets =
+                dataModel.widgetsModel.getAllWidgetComponentsWithoutShortcuts();
 
         List<WidgetItem> servicePredictedItems = new ArrayList<>();
-        List<WidgetItem> localFilteredWidgets = new ArrayList<>();
 
         for (AppTarget app : mTargets) {
-            PackageUserKey packageUserKey = new PackageUserKey(app.getPackageName(), app.getUser());
-            List<WidgetItem> widgets = allWidgets.get(packageUserKey);
-            if (widgets == null || widgets.isEmpty()) {
+            ComponentKey componentKey = new ComponentKey(
+                    new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
+            WidgetItem widget = allWidgets.get(componentKey);
+            if (widget == null) {
                 continue;
             }
             String className = app.getClassName();
             if (!TextUtils.isEmpty(className)) {
-                WidgetItem item = widgets.stream()
-                        .filter(w -> className.equals(w.componentName.getClassName()))
-                        .filter(notOnWorkspace)
-                        .findFirst()
-                        .orElse(null);
-                if (item != null) {
-                    servicePredictedItems.add(item);
-                    continue;
+                if (notOnWorkspace.test(widget)) {
+                    servicePredictedItems.add(widget);
                 }
             }
-            // No widget was added by the service, try local filtering
-            widgets.stream().filter(notOnWorkspace).findFirst()
-                    .ifPresent(localFilteredWidgets::add);
-        }
-        if (servicePredictedItems.isEmpty()) {
-            servicePredictedItems.addAll(localFilteredWidgets);
         }
 
         List<ItemInfo> items;
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index 7eeea84..85ea5fd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -719,6 +719,8 @@
                     } else if (i == MAX_VISIBLE_BUBBLES_COLLAPSED - 1
                             && bubbleCount == MAX_VISIBLE_BUBBLES_COLLAPSED) {
                         bv.setAlpha(0);
+                    } else {
+                        bv.setAlpha(1);
                     }
                 }
             }
diff --git a/quickstep/src/com/android/quickstep/BaseContainerInterface.java b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
index 9955183..b1d511c 100644
--- a/quickstep/src/com/android/quickstep/BaseContainerInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
@@ -41,7 +41,6 @@
 import com.android.launcher3.statehandlers.DesktopVisibilityController;
 import com.android.launcher3.statemanager.BaseState;
 import com.android.launcher3.taskbar.TaskbarUIController;
-import com.android.launcher3.touch.PagedOrientationHandler;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.views.ScrimView;
 import com.android.quickstep.orientation.RecentsPagedOrientationHandler;
@@ -207,10 +206,10 @@
     }
 
     public final void calculateTaskSize(Context context, DeviceProfile dp, Rect outRect,
-            PagedOrientationHandler orientedState) {
+            RecentsPagedOrientationHandler orientationHandler) {
         if (dp.isTablet) {
             if (Flags.enableGridOnlyOverview()) {
-                calculateGridTaskSize(context, dp, outRect, orientedState);
+                calculateGridTaskSize(context, dp, outRect, orientationHandler);
             } else {
                 calculateFocusTaskSize(context, dp, outRect);
             }
@@ -218,15 +217,19 @@
             Resources res = context.getResources();
             float maxScale = res.getFloat(R.dimen.overview_max_scale);
             int taskMargin = dp.overviewTaskMarginPx;
+            // In fake orientation, OverviewActions is hidden and we only leave a margin there.
+            int overviewActionsClaimedSpace = orientationHandler.isLayoutNaturalToLauncher()
+                    ? dp.getOverviewActionsClaimedSpace() : dp.overviewActionsTopMarginPx;
             calculateTaskSizeInternal(
                     context,
                     dp,
                     dp.overviewTaskThumbnailTopMarginPx,
-                    dp.getOverviewActionsClaimedSpace(),
+                    overviewActionsClaimedSpace,
                     res.getDimensionPixelSize(R.dimen.overview_minimum_next_prev_size) + taskMargin,
                     maxScale,
                     Gravity.CENTER,
-                    outRect);
+                    outRect,
+                    orientationHandler);
         }
     }
 
@@ -234,7 +237,7 @@
      * Calculates the taskView size for carousel during app to overview animation on tablets.
      */
     public final void calculateCarouselTaskSize(Context context, DeviceProfile dp, Rect outRect,
-            PagedOrientationHandler orientedState) {
+            RecentsPagedOrientationHandler orientationHandler) {
         if (dp.isTablet && dp.isGestureMode) {
             Resources res = context.getResources();
             float minScale = res.getFloat(R.dimen.overview_carousel_min_scale);
@@ -243,7 +246,7 @@
             calculateTaskSizeInternal(context, dp, gridRect, minScale, Gravity.CENTER | Gravity.TOP,
                     outRect);
         } else {
-            calculateTaskSize(context, dp, outRect, orientedState);
+            calculateTaskSize(context, dp, outRect, orientationHandler);
         }
     }
 
@@ -257,16 +260,39 @@
 
     private void calculateTaskSizeInternal(Context context, DeviceProfile dp, int claimedSpaceAbove,
             int claimedSpaceBelow, int minimumHorizontalPadding, float maxScale, int gravity,
-            Rect outRect) {
-        Rect insets = dp.getInsets();
-
+            Rect outRect, RecentsPagedOrientationHandler orientationHandler) {
         Rect potentialTaskRect = new Rect(0, 0, dp.widthPx, dp.heightPx);
-        potentialTaskRect.inset(insets.left, insets.top, insets.right, insets.bottom);
-        potentialTaskRect.inset(
+
+        Rect insets;
+        if (orientationHandler.isLayoutNaturalToLauncher()) {
+            insets = dp.getInsets();
+        } else {
+            Rect portraitInsets = dp.getInsets();
+            DisplayController displayController = DisplayController.INSTANCE.get(context);
+            Rect deviceRotationInsets = displayController.getInfo().getCurrentBounds().get(
+                    orientationHandler.getRotation()).insets;
+            // Obtain the landscape/seascape insets, and rotate it to portrait perspective.
+            orientationHandler.rotateInsets(deviceRotationInsets, outRect);
+            // Then combine with portrait's insets to leave space for status bar/nav bar in
+            // either orientations.
+            outRect.set(
+                    Math.max(outRect.left, portraitInsets.left),
+                    Math.max(outRect.top, portraitInsets.top),
+                    Math.max(outRect.right, portraitInsets.right),
+                    Math.max(outRect.bottom, portraitInsets.bottom)
+            );
+            insets = outRect;
+        }
+        potentialTaskRect.inset(insets);
+
+        outRect.set(
                 minimumHorizontalPadding,
                 claimedSpaceAbove,
                 minimumHorizontalPadding,
                 claimedSpaceBelow);
+        // Rotate the paddings to portrait perspective,
+        orientationHandler.rotateInsets(outRect, outRect);
+        potentialTaskRect.inset(outRect);
 
         calculateTaskSizeInternal(context, dp, potentialTaskRect, maxScale, gravity, outRect);
     }
@@ -326,7 +352,7 @@
      * Calculates the overview grid non-focused task size for the provided device configuration.
      */
     public final void calculateGridTaskSize(Context context, DeviceProfile dp, Rect outRect,
-            PagedOrientationHandler orientedState) {
+            RecentsPagedOrientationHandler orientationHandler) {
         Resources res = context.getResources();
         Rect potentialTaskRect = new Rect();
         if (Flags.enableGridOnlyOverview()) {
@@ -344,7 +370,7 @@
         int outHeight = Math.round(scale * taskDimension.y);
 
         int gravity = Gravity.TOP;
-        gravity |= orientedState.getRecentsRtlSetting(res) ? Gravity.RIGHT : Gravity.LEFT;
+        gravity |= orientationHandler.getRecentsRtlSetting(res) ? Gravity.RIGHT : Gravity.LEFT;
         Gravity.apply(gravity, outWidth, outHeight, potentialTaskRect, outRect);
     }
 
@@ -352,8 +378,8 @@
      * Calculates the modal taskView size for the provided device configuration
      */
     public final void calculateModalTaskSize(Context context, DeviceProfile dp, Rect outRect,
-            PagedOrientationHandler orientedState) {
-        calculateTaskSize(context, dp, outRect, orientedState);
+            RecentsPagedOrientationHandler orientationHandler) {
+        calculateTaskSize(context, dp, outRect, orientationHandler);
         boolean isGridOnlyOverview = dp.isTablet && Flags.enableGridOnlyOverview();
         int claimedSpaceBelow = isGridOnlyOverview
                 ? dp.overviewActionsTopMarginPx + dp.overviewActionsHeight + dp.stashedTaskbarHeight
@@ -372,6 +398,7 @@
                 minimumHorizontalPadding,
                 1f /*maxScale*/,
                 Gravity.CENTER_HORIZONTAL | Gravity.BOTTOM,
-                outRect);
+                outRect,
+                orientationHandler);
     }
 }
diff --git a/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.kt b/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.kt
index 8f8cc6e..ec04cb7 100644
--- a/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.kt
+++ b/quickstep/src/com/android/quickstep/orientation/LandscapePagedViewHandler.kt
@@ -121,6 +121,10 @@
 
     override fun getEnd(rect: RectF): Float = rect.bottom
 
+    override fun rotateInsets(insets: Rect, outInsets: Rect) {
+        outInsets.set(insets.bottom, insets.left, insets.top, insets.right)
+    }
+
     override fun getClearAllSidePadding(view: View, isRtl: Boolean): Int =
         if (isRtl) view.paddingBottom / 2 else -view.paddingTop / 2
 
diff --git a/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java b/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java
index f6284d5..758a737 100644
--- a/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java
+++ b/quickstep/src/com/android/quickstep/orientation/PortraitPagedViewHandler.java
@@ -130,6 +130,11 @@
     }
 
     @Override
+    public void rotateInsets(@NonNull Rect insets, @NonNull Rect outInsets) {
+        outInsets.set(insets);
+    }
+
+    @Override
     public int getClearAllSidePadding(View view, boolean isRtl) {
         return (isRtl ? view.getPaddingRight() : - view.getPaddingLeft()) / 2;
     }
diff --git a/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.kt b/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.kt
index 5bc1be8..df4b030 100644
--- a/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.kt
+++ b/quickstep/src/com/android/quickstep/orientation/RecentsPagedOrientationHandler.kt
@@ -70,6 +70,9 @@
 
     fun getEnd(rect: RectF): Float
 
+    /** Rotate the provided insets to portrait perspective. */
+    fun rotateInsets(insets: Rect, outInsets: Rect)
+
     fun getClearAllSidePadding(view: View, isRtl: Boolean): Int
 
     fun getSecondaryDimension(view: View): Int
diff --git a/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.kt b/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.kt
index 46e4b0c..333359f 100644
--- a/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.kt
+++ b/quickstep/src/com/android/quickstep/orientation/SeascapePagedViewHandler.kt
@@ -44,6 +44,10 @@
 import com.android.quickstep.views.IconAppChipView
 
 class SeascapePagedViewHandler : LandscapePagedViewHandler() {
+    override fun rotateInsets(insets: Rect, outInsets: Rect) {
+        outInsets.set(insets.top, insets.right, insets.bottom, insets.left)
+    }
+
     override val secondaryTranslationDirectionFactor: Int = -1
 
     override fun getSplitTranslationDirectionFactor(
diff --git a/quickstep/src/com/android/quickstep/views/FloatingTaskView.java b/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
index 0e451b5..acbb2ec 100644
--- a/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
@@ -187,7 +187,13 @@
                     viewBounds, false /* ignoreTransform */, null /* recycle */,
                     mStartingPosition);
         }
-
+        // In some cases originalView is off-screen so we don't get a valid starting position
+        // ex. on rotation
+        // TODO(b/345556328) We shouldn't be animating if starting position of view isn't ready
+        if (mStartingPosition.isEmpty()) {
+            // Set to non empty for now so calculations in #update() don't break
+            mStartingPosition.set(0, 0, 1, 1);
+        }
         final BaseDragLayer.LayoutParams lp = new BaseDragLayer.LayoutParams(
                 Math.round(mStartingPosition.width()),
                 Math.round(mStartingPosition.height()));
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 4fd686d..5e79743 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -825,7 +825,10 @@
         taskContainers.forEach {
             val thumbnailBounds = Rect()
             if (relativeToDragLayer) {
-                container.dragLayer.getDescendantRectRelativeToSelf(it.snapshotView, bounds)
+                container.dragLayer.getDescendantRectRelativeToSelf(
+                    it.snapshotView,
+                    thumbnailBounds
+                )
             } else {
                 thumbnailBounds.set(it.snapshotView)
             }
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 8702f70..7b57c81 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -184,7 +184,7 @@
     }
 
     @Test
-    public void widgetsRecommendationRan_shouldReturnPackageWidgetsWhenEmpty() {
+    public void widgetsRecommendationRan_shouldReturnEmptyWidgetsWhenEmpty() {
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
 
             // Not installed widget
@@ -204,19 +204,12 @@
                     newWidgetsPredicationTask(List.of(widget5, widget3, widget4, widget1)));
             runOnExecutorSync(MAIN_EXECUTOR, () -> { });
 
-            // THEN only 2 widgets are returned because the launcher only filters out
-            // non-exist widgets.
+            // Only widgets suggested by prediction system are returned.
             List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
                     .stream()
                     .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
                     .collect(Collectors.toList());
-            assertThat(recommendedWidgets).hasSize(2);
-            recommendedWidgets.forEach(pendingAddWidgetInfo ->
-                    assertThat(pendingAddWidgetInfo.recommendationCategory).isNotNull()
-            );
-            // Another widget from the same package
-            assertWidgetInfo(recommendedWidgets.get(0).info, mApp4Provider2);
-            assertWidgetInfo(recommendedWidgets.get(1).info, mApp1Provider1);
+            assertThat(recommendedWidgets).hasSize(0);
         });
     }
 
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 5e0edb3..4cba0b5 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -32,6 +32,7 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
+import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
@@ -130,6 +131,22 @@
     }
 
     /**
+     * Returns a map of widget component keys to corresponding widget items. Excludes the
+     * shortcuts.
+     */
+    public synchronized Map<ComponentKey, WidgetItem> getAllWidgetComponentsWithoutShortcuts() {
+        if (!WIDGETS_ENABLED) {
+            return Collections.emptyMap();
+        }
+        Map<ComponentKey, WidgetItem> widgetsMap = new HashMap<>();
+        mWidgetsList.forEach((packageItemInfo, widgetsAndShortcuts) ->
+                widgetsAndShortcuts.stream().filter(item -> item.widgetInfo != null).forEach(
+                        item -> widgetsMap.put(new ComponentKey(item.componentName, item.user),
+                                item)));
+        return widgetsMap;
+    }
+
+    /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
      */
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 21eee55..96b3457 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -405,7 +405,7 @@
             navigationMode = wmProxy.getNavigationMode(displayInfoContext);
 
             mPerDisplayBounds.putAll(perDisplayBoundsCache);
-            List<WindowBounds> cachedValue = mPerDisplayBounds.get(normalizedDisplayInfo);
+            List<WindowBounds> cachedValue = getCurrentBounds();
 
             realBounds = wmProxy.getRealBounds(displayInfoContext, displayInfo);
             if (cachedValue == null) {
@@ -415,7 +415,7 @@
                 FileLog.e(TAG, "(Invalid Cache) perDisplayBounds : " + mPerDisplayBounds);
                 mPerDisplayBounds.clear();
                 mPerDisplayBounds.putAll(wmProxy.estimateInternalDisplayBounds(displayInfoContext));
-                cachedValue = mPerDisplayBounds.get(normalizedDisplayInfo);
+                cachedValue = getCurrentBounds();
                 if (cachedValue == null) {
                     FileLog.e(TAG, "normalizedDisplayInfo not found in estimation: "
                             + normalizedDisplayInfo);
@@ -505,6 +505,13 @@
             return Collections.unmodifiableSet(mPerDisplayBounds.keySet());
         }
 
+        /**
+         * Returns all {@link WindowBounds}s for the current display.
+         */
+        public List<WindowBounds> getCurrentBounds() {
+            return mPerDisplayBounds.get(normalizedDisplayInfo);
+        }
+
         public int getDensityDpi() {
             return densityDpi;
         }
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index 2bb485a..35372d3 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -372,6 +372,13 @@
         }
     }
 
+    /**
+     * Returns a view (holding the previews) that can be dragged and dropped.
+     */
+    public View getDragAndDropView() {
+        return mWidgetImageContainer;
+    }
+
     public WidgetImageView getWidgetView() {
         return mWidgetImage;
     }
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/UserBadgeDrawableTest.kt b/tests/multivalentTests/src/com/android/launcher3/icons/UserBadgeDrawableTest.kt
new file mode 100644
index 0000000..d611ae8
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/UserBadgeDrawableTest.kt
@@ -0,0 +1,76 @@
+/*
+ * 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.launcher3.icons
+
+import android.graphics.Canvas
+import android.graphics.Color
+import android.graphics.Paint
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.icons.UserBadgeDrawable.SHADOW_COLOR
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [UserBadgeDrawable] */
+@RunWith(AndroidJUnit4::class)
+class UserBadgeDrawableTest {
+    private val context = InstrumentationRegistry.getInstrumentation().targetContext
+    private val canvas = mock<Canvas>()
+    private val systemUnderTest =
+        UserBadgeDrawable(context, R.drawable.ic_work_app_badge, R.color.badge_tint_work, false)
+
+    @Test
+    fun draw_opaque() {
+        val colorList = mutableListOf<Int>()
+        whenever(
+            canvas.drawCircle(
+                any(),
+                any(),
+                any(),
+                any()
+            )
+        ).then { colorList.add(it.getArgument<Paint>(3).color) }
+
+        systemUnderTest.alpha = 255
+        systemUnderTest.draw(canvas)
+
+        assertThat(colorList).containsExactly(SHADOW_COLOR, Color.WHITE)
+    }
+
+    @Test
+    fun draw_transparent() {
+        val colorList = mutableListOf<Int>()
+        whenever(
+            canvas.drawCircle(
+                any(),
+                any(),
+                any(),
+                any()
+            )
+        ).then { colorList.add(it.getArgument<Paint>(3).color) }
+
+        systemUnderTest.alpha = 0
+        systemUnderTest.draw(canvas)
+
+        assertThat(colorList).hasSize(2)
+        assertThat(Color.valueOf(colorList[0]).alpha()).isEqualTo(0)
+        assertThat(Color.valueOf(colorList[1]).alpha()).isEqualTo(0)
+    }
+}