Merge "Fix bubble bar and handle visibility on first bubble" into main
diff --git a/aconfig/launcher_search.aconfig b/aconfig/launcher_search.aconfig
index b243922..b98eee6 100644
--- a/aconfig/launcher_search.aconfig
+++ b/aconfig/launcher_search.aconfig
@@ -45,8 +45,11 @@
 
 
 flag {
-    name: "private_space_floating_mask_view"
+    name: "private_space_add_floating_mask_view"
     namespace: "launcher_search"
     description: "This flag enables the floating mask view as part of the Private Space animation. "
     bug: "339850589"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissController.java
index 0e6fa3c..a6096e2 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissController.java
@@ -169,7 +169,8 @@
 
     private void setupMagnetizedObject(@NonNull View magnetizedView) {
         mMagnetizedObject = new MagnetizedObject<>(mActivity.getApplicationContext(),
-                magnetizedView, DynamicAnimation.TRANSLATION_X, DynamicAnimation.TRANSLATION_Y) {
+                magnetizedView, BubbleDragController.DRAG_TRANSLATION_X,
+                DynamicAnimation.TRANSLATION_Y) {
             @Override
             public float getWidth(@NonNull View underlyingObject) {
                 return underlyingObject.getWidth() * underlyingObject.getScaleX();
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragAnimator.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragAnimator.java
index 287e906..7aed2d2 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragAnimator.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragAnimator.java
@@ -60,7 +60,6 @@
     private final float mBubbleFocusedScale;
     private final float mBubbleCapturedScale;
     private final float mDismissCapturedScale;
-    private final FloatPropertyCompat<View> mTranslationXProperty;
 
     /**
      * Should be initialised for each dragged view
@@ -82,28 +81,9 @@
         if (view instanceof BubbleBarView) {
             mBubbleFocusedScale = SCALE_BUBBLE_BAR_FOCUSED;
             mBubbleCapturedScale = mDismissCapturedScale;
-            mTranslationXProperty = DynamicAnimation.TRANSLATION_X;
         } else {
             mBubbleFocusedScale = SCALE_BUBBLE_FOCUSED;
             mBubbleCapturedScale = SCALE_BUBBLE_CAPTURED;
-            // Wrap BubbleView.DRAG_TRANSLATION_X as it can't be cast to FloatPropertyCompat<View>
-            mTranslationXProperty = new FloatPropertyCompat<>(
-                    BubbleView.DRAG_TRANSLATION_X.getName()) {
-                @Override
-                public float getValue(View object) {
-                    if (object instanceof BubbleView bubbleView) {
-                        return BubbleView.DRAG_TRANSLATION_X.get(bubbleView);
-                    }
-                    return 0;
-                }
-
-                @Override
-                public void setValue(View object, float value) {
-                    if (object instanceof BubbleView bubbleView) {
-                        BubbleView.DRAG_TRANSLATION_X.setValue(bubbleView, value);
-                    }
-                }
-            };
         }
     }
 
@@ -140,7 +120,7 @@
         mBubbleAnimator
                 .spring(DynamicAnimation.SCALE_X, 1f)
                 .spring(DynamicAnimation.SCALE_Y, 1f)
-                .spring(mTranslationXProperty, restingPosition.x, velocity.x,
+                .spring(BubbleDragController.DRAG_TRANSLATION_X, restingPosition.x, velocity.x,
                         mTranslationConfig)
                 .spring(DynamicAnimation.TRANSLATION_Y, restingPosition.y, velocity.y,
                         mTranslationConfig)
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java
index e04c1b1..fbd1b88 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java
@@ -24,6 +24,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.dynamicanimation.animation.FloatPropertyCompat;
 
 import com.android.launcher3.taskbar.TaskbarActivityContext;
 import com.android.wm.shell.common.bubbles.BaseBubblePinController.LocationChangeListener;
@@ -38,6 +39,37 @@
  * Restores initial position of dragged view if released outside of the dismiss target.
  */
 public class BubbleDragController {
+
+    /**
+     * Property to update dragged bubble x-translation value.
+     * <p>
+     * When applied to {@link BubbleView}, will use set the translation through
+     * {@link BubbleView#getDragTranslationX()} and {@link BubbleView#setDragTranslationX(float)}
+     * methods.
+     * <p>
+     * When applied to {@link BubbleBarView}, will use {@link View#getTranslationX()} and
+     * {@link View#setTranslationX(float)}.
+     */
+    public static final FloatPropertyCompat<View> DRAG_TRANSLATION_X = new FloatPropertyCompat<>(
+            "dragTranslationX") {
+        @Override
+        public float getValue(View view) {
+            if (view instanceof BubbleView bubbleView) {
+                return bubbleView.getDragTranslationX();
+            }
+            return view.getTranslationX();
+        }
+
+        @Override
+        public void setValue(View view, float value) {
+            if (view instanceof BubbleView bubbleView) {
+                bubbleView.setDragTranslationX(value);
+            } else {
+                view.setTranslationX(value);
+            }
+        }
+    };
+
     private final TaskbarActivityContext mActivity;
     private BubbleBarController mBubbleBarController;
     private BubbleBarViewController mBubbleBarViewController;
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
index 61a6bce..2e37dc7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
@@ -22,13 +22,11 @@
 import android.graphics.Outline;
 import android.graphics.Rect;
 import android.util.AttributeSet;
-import android.util.FloatProperty;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewOutlineProvider;
 import android.widget.ImageView;
 
-import androidx.annotation.NonNull;
 import androidx.constraintlayout.widget.ConstraintLayout;
 
 import com.android.launcher3.R;
@@ -49,25 +47,6 @@
     public static final int DEFAULT_PATH_SIZE = 100;
 
     /**
-     * Property to update drag translation value.
-     *
-     * @see BubbleView#getDragTranslationX()
-     * @see BubbleView#setDragTranslationX(float)
-     */
-    public static final FloatProperty<BubbleView> DRAG_TRANSLATION_X = new FloatProperty<>(
-            "dragTranslationX") {
-        @Override
-        public void setValue(@NonNull BubbleView bubbleView, float value) {
-            bubbleView.setDragTranslationX(value);
-        }
-
-        @Override
-        public Float get(BubbleView bubbleView) {
-            return bubbleView.getDragTranslationX();
-        }
-    };
-
-    /**
      * Flags that suppress the visibility of the 'new' dot or the app badge, for one reason or
      * another. If any of these flags are set, the dot will not be shown.
      * If {@link SuppressionFlag#BEHIND_STACK} then the app badge will not be shown.
diff --git a/res/layout/private_space_header.xml b/res/layout/private_space_header.xml
index cefe394..9c0f129 100644
--- a/res/layout/private_space_header.xml
+++ b/res/layout/private_space_header.xml
@@ -37,7 +37,7 @@
         android:gravity="center_vertical"
         android:layout_alignParentEnd="true"
         android:animateLayoutChanges="false">
-        <ImageButton
+        <com.android.launcher3.allapps.PrivateSpaceSettingsButton
             android:id="@+id/ps_settings_button"
             android:layout_width="@dimen/ps_header_image_height"
             android:layout_height="@dimen/ps_header_image_height"
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 7f72526..6d2c8fa 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -163,6 +163,7 @@
 import androidx.annotation.StringRes;
 import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
+import androidx.core.os.BuildCompat;
 import androidx.window.embedding.RuleController;
 
 import com.android.launcher3.DropTarget.DragObject;
@@ -589,7 +590,8 @@
         setTitle(R.string.home_screen);
         mStartupLatencyLogger.logEnd(LAUNCHER_LATENCY_STARTUP_ACTIVITY_ON_CREATE);
 
-        if (com.android.launcher3.Flags.enableTwoPaneLauncherSettings()) {
+        if (BuildCompat.isAtLeastV()
+                && com.android.launcher3.Flags.enableTwoPaneLauncherSettings()) {
             RuleController.getInstance(this).setRules(
                     RuleController.parseRules(this, R.xml.split_configuration));
         }
@@ -2789,7 +2791,7 @@
     }
 
     private void updateDisallowBack() {
-        if (Flags.enableDesktopWindowingMode()) {
+        if (BuildCompat.isAtLeastV() && Flags.enableDesktopWindowingMode()) {
             // TODO(b/330183377) disable back in launcher when when we productionize
             return;
         }
@@ -3146,4 +3148,4 @@
     }
 
     // End of Getters and Setters
-}
\ No newline at end of file
+}
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index a4ae1c8..3b8ff62 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -39,6 +39,7 @@
 import android.util.Log;
 
 import androidx.annotation.Nullable;
+import androidx.core.os.BuildCompat;
 
 import com.android.launcher3.graphics.IconShape;
 import com.android.launcher3.icons.IconCache;
@@ -107,7 +108,7 @@
         mOnTerminateCallback.add(() ->
                 mContext.getSystemService(LauncherApps.class).unregisterCallback(callbacks));
 
-        if (Flags.enableSupportForArchiving()) {
+        if (BuildCompat.isAtLeastV() && Flags.enableSupportForArchiving()) {
             ArchiveCompatibilityParams params = new ArchiveCompatibilityParams();
             params.setEnableUnarchivalConfirmation(false);
             launcherApps.setArchiveCompatibility(params);
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 2765c5c..32b5cfa 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -30,7 +30,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_LOCK_ANIMATION_BEGIN;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_LOCK_ANIMATION_END;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_LOCK_TAP;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_SETTINGS_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_UNLOCK_ANIMATION_BEGIN;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_UNLOCK_ANIMATION_END;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_UNLOCK_TAP;
@@ -52,7 +51,6 @@
 import android.os.UserManager;
 import android.view.View;
 import android.view.ViewGroup;
-import android.widget.ImageButton;
 import android.widget.ImageView;
 import android.widget.RelativeLayout;
 import android.widget.TextView;
@@ -229,30 +227,6 @@
         resetPrivateSpaceDecorator(updatedState);
     }
 
-    /**
-     * Opens the Private Space Settings Page.
-     *
-     * @param view the view that was clicked to open the settings page and which will be the same
-     *             view to animate back. Otherwise if there is no view, simply start the activity.
-     */
-    public void openPrivateSpaceSettings(View view) {
-        if (mPrivateSpaceSettingsAvailable) {
-            Context context = mAllApps.getContext();
-            Intent intent = ApiWrapper.INSTANCE.get(context).getPrivateSpaceSettingsIntent();
-            if (view == null) {
-                context.startActivity(intent);
-                return;
-            }
-            ActivityContext activityContext = ActivityContext.lookupContext(context);
-            AppInfo itemInfo = new AppInfo();
-            itemInfo.id = CONTAINER_PRIVATESPACE;
-            itemInfo.componentName = intent.getComponent();
-            itemInfo.container = CONTAINER_PRIVATESPACE;
-            view.setTag(itemInfo);
-            activityContext.startActivitySafely(view, intent, itemInfo);
-        }
-    }
-
     /** Returns whether or not Private Space Settings Page is available. */
     public boolean isPrivateSpaceSettingsAvailable() {
         return mPrivateSpaceSettingsAvailable;
@@ -385,7 +359,7 @@
         updateHeaderOnClickListener(mPSHeader);
 
         //Add image and action for private space settings button
-        ImageButton settingsButton = mPSHeader.findViewById(R.id.ps_settings_button);
+        PrivateSpaceSettingsButton settingsButton = mPSHeader.findViewById(R.id.ps_settings_button);
         assert settingsButton != null;
         updatePrivateSpaceSettingsButton(settingsButton);
 
@@ -439,15 +413,10 @@
         setQuietMode(lock);
     }
 
-    private void updatePrivateSpaceSettingsButton(ImageButton settingsButton) {
+    private void updatePrivateSpaceSettingsButton(PrivateSpaceSettingsButton settingsButton) {
         if (getCurrentState() == STATE_ENABLED
                 && isPrivateSpaceSettingsAvailable()) {
             settingsButton.setVisibility(VISIBLE);
-            settingsButton.setOnClickListener(
-                    view -> {
-                        logEvents(LAUNCHER_PRIVATE_SPACE_SETTINGS_TAP);
-                        openPrivateSpaceSettings(view);
-                    });
         } else {
             settingsButton.setVisibility(GONE);
         }
@@ -723,7 +692,7 @@
 
     /** Fades out the private space container. */
     private ValueAnimator translateFloatingMaskView(boolean animateIn) {
-        if (!Flags.privateSpaceFloatingMaskView() || mFloatingMaskView == null) {
+        if (!Flags.privateSpaceAddFloatingMaskView() || mFloatingMaskView == null) {
             return new ValueAnimator();
         }
         // Translate base on the height amount. Translates out on expand and in on collapse.
@@ -834,7 +803,7 @@
     }
 
     private void attachFloatingMaskView(boolean expand) {
-        if (!Flags.privateSpaceFloatingMaskView()) {
+        if (!Flags.privateSpaceAddFloatingMaskView()) {
             return;
         }
         mFloatingMaskView = (FloatingMaskView) mAllApps.getLayoutInflater().inflate(
diff --git a/src/com/android/launcher3/allapps/PrivateSpaceSettingsButton.java b/src/com/android/launcher3/allapps/PrivateSpaceSettingsButton.java
new file mode 100644
index 0000000..43e42ff
--- /dev/null
+++ b/src/com/android/launcher3/allapps/PrivateSpaceSettingsButton.java
@@ -0,0 +1,81 @@
+/*
+ * 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.allapps;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PRIVATESPACE;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_SETTINGS_TAP;
+
+import android.content.Context;
+import android.content.Intent;
+import android.util.AttributeSet;
+import android.view.View;
+import android.widget.ImageButton;
+
+import com.android.launcher3.logging.StatsLogManager;
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.util.ApiWrapper;
+import com.android.launcher3.views.ActivityContext;
+
+public class PrivateSpaceSettingsButton extends ImageButton implements View.OnClickListener {
+
+    private final ActivityContext mActivityContext;
+    private final StatsLogManager mStatsLogManager;
+    private final Intent mPrivateSpaceSettingsIntent;
+
+    public PrivateSpaceSettingsButton(Context context) {
+        this(context, null, 0);
+    }
+
+    public PrivateSpaceSettingsButton(Context context, AttributeSet attrs) {
+        this(context, attrs, 0);
+    }
+
+    public PrivateSpaceSettingsButton(Context context, AttributeSet attrs, int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+        mActivityContext = ActivityContext.lookupContext(context);
+        mStatsLogManager = mActivityContext.getStatsLogManager();
+        mPrivateSpaceSettingsIntent =
+                ApiWrapper.INSTANCE.get(context).getPrivateSpaceSettingsIntent();
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        setOnClickListener(this);
+    }
+
+    @Override
+    public void onClick(View view) {
+        mStatsLogManager.logger().log(LAUNCHER_PRIVATE_SPACE_SETTINGS_TAP);
+        AppInfo privateSpaceSettingsItemInfo = createPrivateSpaceSettingsAppInfo();
+        view.setTag(privateSpaceSettingsItemInfo);
+        mActivityContext.startActivitySafely(
+                view,
+                mPrivateSpaceSettingsIntent,
+                privateSpaceSettingsItemInfo);
+    }
+
+    AppInfo createPrivateSpaceSettingsAppInfo() {
+        AppInfo itemInfo = new AppInfo();
+        itemInfo.id = CONTAINER_PRIVATESPACE;
+        if (mPrivateSpaceSettingsIntent != null) {
+            itemInfo.componentName = mPrivateSpaceSettingsIntent.getComponent();
+        }
+        itemInfo.container = CONTAINER_PRIVATESPACE;
+        return itemInfo;
+    }
+}
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index ec0a222..85eb39b 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -91,7 +91,8 @@
  */
 @TargetApi(Build.VERSION_CODES.O)
 public class AddItemActivity extends BaseActivity
-        implements OnLongClickListener, OnTouchListener, AbstractSlideInView.OnCloseListener {
+        implements OnLongClickListener, OnTouchListener, AbstractSlideInView.OnCloseListener,
+        WidgetCell.PreviewReadyListener {
 
     private static final int SHADOW_SIZE = 10;
 
@@ -142,6 +143,7 @@
         mDragLayer = findViewById(R.id.add_item_drag_layer);
         mDragLayer.recreateControllers();
         mWidgetCell = findViewById(R.id.widget_cell);
+        mWidgetCell.addPreviewReadyListener(this);
         mAccessibilityManager =
                 getApplicationContext().getSystemService(AccessibilityManager.class);
 
@@ -454,4 +456,11 @@
                 .withItemInfo((ItemInfo) mWidgetCell.getWidgetView().getTag())
                 .log(command);
     }
+
+    @Override
+    public void onPreviewAvailable() {
+        // Set the preview height based on "the only" widget's preview.
+        mWidgetCell.setParentAlignedPreviewHeight(mWidgetCell.getPreviewContentHeight());
+        mWidgetCell.post(mWidgetCell::requestLayout);
+    }
 }
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index eac2ce7..2bb485a 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -103,6 +103,8 @@
     private Size mWidgetSize;
 
     private final DatabaseWidgetPreviewLoader mWidgetPreviewLoader;
+    @Nullable
+    private PreviewReadyListener mPreviewReadyListener = null;
 
     protected CancellableTask mActiveRequest;
     private boolean mAnimatePreview = true;
@@ -118,7 +120,8 @@
 
     private CancellableTask mIconLoadRequest;
     private boolean mIsShowingAddButton = false;
-
+    // Height enforced by the parent to align all widget cells displayed by it.
+    private int mParentAlignedPreviewHeight;
     public WidgetCell(Context context) {
         this(context, null);
     }
@@ -190,6 +193,8 @@
         mWidgetDims.setText(null);
         mWidgetDescription.setText(null);
         mWidgetDescription.setVisibility(GONE);
+        mPreviewReadyListener = null;
+        mParentAlignedPreviewHeight = 0;
         showDescription(true);
         showDimensions(true);
 
@@ -338,8 +343,8 @@
 
     private void updateAppWidgetHostScale(NavigableAppWidgetHostView view) {
         // Scale the content such that all of the content is visible
-        int contentWidth = view.getWidth();
-        int contentHeight = view.getHeight();
+        float contentWidth = view.getWidth();
+        float contentHeight = view.getHeight();
 
         if (view.getChildCount() == 1) {
             View content = view.getChildAt(0);
@@ -359,6 +364,12 @@
             mAppWidgetHostViewScale = Math.min(pWidth / contentWidth, pHeight / contentHeight);
         }
         view.setScaleToFit(mAppWidgetHostViewScale);
+
+        // layout based previews maybe ready at this point to inspect their inner height.
+        if (mPreviewReadyListener != null) {
+            mPreviewReadyListener.onPreviewAvailable();
+            mPreviewReadyListener = null;
+        }
     }
 
     public WidgetImageView getWidgetView() {
@@ -384,6 +395,12 @@
                 removeView(mAppWidgetHostViewPreview);
                 mAppWidgetHostViewPreview = null;
             }
+
+            // Drawables of the image previews are available at this point to measure.
+            if (mPreviewReadyListener != null) {
+                mPreviewReadyListener.onPreviewAvailable();
+                mPreviewReadyListener = null;
+            }
         }
 
         if (mAnimatePreview) {
@@ -489,14 +506,20 @@
         // mPreviewContainerScale ensures the needed scaling with respect to original widget size.
         mAppWidgetHostViewScale = mPreviewContainerScale;
         containerLp.width = mPreviewContainerSize.getWidth();
-        containerLp.height = mPreviewContainerSize.getHeight();
+        int height = mPreviewContainerSize.getHeight();
 
         // If we don't have enough available width, scale the preview container to fit.
         if (containerLp.width > maxWidth) {
             containerLp.width = maxWidth;
             mAppWidgetHostViewScale = (float) containerLp.width / mPreviewContainerSize.getWidth();
-            containerLp.height = Math.round(
-                    mPreviewContainerSize.getHeight() * mAppWidgetHostViewScale);
+            height = Math.round(mPreviewContainerSize.getHeight() * mAppWidgetHostViewScale);
+        }
+
+        // Use parent aligned height in set.
+        if (mParentAlignedPreviewHeight > 0) {
+            containerLp.height = Math.min(height, mParentAlignedPreviewHeight);
+        } else {
+            containerLp.height = height;
         }
 
         // No need to call mWidgetImageContainer.setLayoutParams as we are in measure pass
@@ -513,6 +536,42 @@
     }
 
     /**
+     * Sets the height of the preview as adjusted by the parent to have this cell's content aligned
+     * with other cells displayed by the parent.
+     */
+    public void setParentAlignedPreviewHeight(int previewHeight) {
+        mParentAlignedPreviewHeight = previewHeight;
+    }
+
+    /**
+     * Returns the height of the preview without any empty space.
+     * In case of appwidget host views, it returns the height of first child. This way, if preview
+     * view provided by an app doesn't fill bounds, this will return actual height without white
+     * space.
+     */
+    public int getPreviewContentHeight() {
+        // By default assume scaled height.
+        int height = Math.round(mPreviewContainerScale * mWidgetSize.getHeight());
+
+        if (mWidgetImage != null && mWidgetImage.getDrawable() != null) {
+            // getBitmapBounds returns the scaled bounds.
+            Rect bitmapBounds = mWidgetImage.getBitmapBounds();
+            height = bitmapBounds.height();
+        } else if (mAppWidgetHostViewPreview != null
+                && mAppWidgetHostViewPreview.getChildCount() == 1) {
+            int contentHeight = Math.round(
+                    mPreviewContainerScale * mWidgetSize.getHeight());
+            int previewInnerHeight = Math.round(
+                    mAppWidgetHostViewScale * mAppWidgetHostViewPreview.getChildAt(
+                            0).getMeasuredHeight());
+            // Use either of the inner scaled height or the scaled widget height
+            height = Math.min(contentHeight, previewInnerHeight);
+        }
+
+        return height;
+    }
+
+    /**
      * Loads a high resolution package icon to show next to the widget title.
      */
     public void loadHighResPackageIcon() {
@@ -651,4 +710,19 @@
         }
         return false;
     }
+
+    /**
+     * Listener to notify when previews are available.
+     */
+    public void addPreviewReadyListener(PreviewReadyListener previewReadyListener) {
+        mPreviewReadyListener = previewReadyListener;
+    }
+
+    /**
+     * Listener interface for subscribers to listen to preview's availability.
+     */
+    public interface PreviewReadyListener {
+        /** Handler on to invoke when previews are available. */
+        void onPreviewAvailable();
+    }
 }
diff --git a/src/com/android/launcher3/widget/WidgetTableRow.java b/src/com/android/launcher3/widget/WidgetTableRow.java
new file mode 100644
index 0000000..a5312ce
--- /dev/null
+++ b/src/com/android/launcher3/widget/WidgetTableRow.java
@@ -0,0 +1,90 @@
+/*
+ * 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.widget;
+
+import android.content.Context;
+import android.util.AttributeSet;
+import android.widget.TableRow;
+
+/**
+ * A row of {@link WidgetCell}s that can be displayed in a table.
+ */
+public class WidgetTableRow extends TableRow implements WidgetCell.PreviewReadyListener {
+    private int mNumOfReadyCells;
+    private int mNumOfCells;
+    private int mResizeDelay;
+
+    public WidgetTableRow(Context context) {
+        super(context);
+    }
+    public WidgetTableRow(Context context, AttributeSet attrs) {
+        super(context, attrs);
+    }
+
+    @Override
+    public void onPreviewAvailable() {
+        mNumOfReadyCells++;
+
+        // Once all previews are loaded, find max visible height and adjust the preview containers.
+        if (mNumOfReadyCells == mNumOfCells) {
+            resize();
+        }
+    }
+
+    private void resize() {
+        int previewHeight = 0;
+        // get the maximum height of each widget preview
+        for (int i = 0; i < getChildCount(); i++) {
+            WidgetCell widgetCell = (WidgetCell) getChildAt(i);
+            previewHeight = Math.max(widgetCell.getPreviewContentHeight(), previewHeight);
+        }
+        if (mResizeDelay > 0) {
+            postDelayed(() -> setAlpha(1f), mResizeDelay);
+        }
+        if (previewHeight > 0) {
+            for (int i = 0; i < getChildCount(); i++) {
+                WidgetCell widgetCell = (WidgetCell) getChildAt(i);
+                widgetCell.setParentAlignedPreviewHeight(previewHeight);
+                widgetCell.postDelayed(widgetCell::requestLayout, mResizeDelay);
+            }
+        }
+    }
+
+    @Override
+    protected void onLayout(boolean changed, int l, int t, int r, int b) {
+        super.onLayout(changed, l, t, r, b);
+    }
+
+    /**
+     * Sets up the row to display the provided number of numOfCells.
+     *
+     * @param numOfCells    number of numOfCells in the row
+     * @param resizeDelayMs time to wait in millis before making any layout size adjustments e.g. we
+     *                      want to wait for list expand collapse animation before resizing the
+     *                      cell previews.
+     */
+    public void setupRow(int numOfCells, int resizeDelayMs) {
+        mNumOfCells = numOfCells;
+        mNumOfReadyCells = 0;
+
+        mResizeDelay = resizeDelayMs;
+        // For delayed resize, reveal contents only after resize is done.
+        if (mResizeDelay > 0) {
+            setAlpha(0);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/WidgetsBottomSheet.java b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
index 4ea2426..894099d 100644
--- a/src/com/android/launcher3/widget/WidgetsBottomSheet.java
+++ b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
@@ -31,7 +31,6 @@
 import android.view.animation.Interpolator;
 import android.widget.ScrollView;
 import android.widget.TableLayout;
-import android.widget.TableRow;
 import android.widget.TextView;
 
 import androidx.annotation.Px;
@@ -137,8 +136,9 @@
                 mActivityContext.getDeviceProfile(), mMaxHorizontalSpan,
                 mWidgetCellHorizontalPadding)
                 .forEach(row -> {
-                    TableRow tableRow = new TableRow(getContext());
+                    WidgetTableRow tableRow = new WidgetTableRow(getContext());
                     tableRow.setGravity(Gravity.TOP);
+                    tableRow.setupRow(row.size(), /*resizeDelayMs=*/ 0);
                     row.forEach(widgetItem -> {
                         WidgetCell widget = addItemCell(tableRow);
                         widget.applyFromCellItem(widgetItem);
@@ -163,9 +163,10 @@
         return super.onControllerInterceptTouchEvent(ev);
     }
 
-    protected WidgetCell addItemCell(ViewGroup parent) {
+    protected WidgetCell addItemCell(WidgetTableRow parent) {
         WidgetCell widget = (WidgetCell) LayoutInflater.from(getContext())
                 .inflate(R.layout.widget_cell, parent, false);
+        widget.addPreviewReadyListener(parent);
         widget.setOnClickListener(this);
 
         View previewContainer = widget.findViewById(R.id.widget_preview_container);
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java
index 4f73e66..9260af9 100644
--- a/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationsView.java
@@ -53,6 +53,7 @@
  */
 public final class WidgetRecommendationsView extends PagedView<PageIndicatorDots> {
     private @Px float mAvailableHeight = Float.MAX_VALUE;
+    private @Px float mAvailableWidth = 0;
     private static final String INITIALLY_DISPLAYED_WIDGETS_STATE_KEY =
             "widgetRecommendationsView:mDisplayedWidgets";
     private static final int MAX_CATEGORIES = 3;
@@ -152,6 +153,7 @@
             final @Px float availableHeight, final @Px int availableWidth,
             final @Px int cellPadding) {
         this.mAvailableHeight = availableHeight;
+        this.mAvailableWidth = availableWidth;
         clear();
 
         Set<ComponentName> displayedWidgets = maybeDisplayInTable(recommendedWidgets,
@@ -187,6 +189,7 @@
             DeviceProfile deviceProfile, final @Px float availableHeight,
             final @Px int availableWidth, final @Px int cellPadding, final int requestedPage) {
         this.mAvailableHeight = availableHeight;
+        this.mAvailableWidth = availableWidth;
         Context context = getContext();
         // For purpose of recommendations section, we don't want paging dots to be halved in two
         // pane display, so, we always provide isTwoPanels = "false".
@@ -280,17 +283,24 @@
         boolean hasMultiplePages = getChildCount() > 0;
 
         if (hasMultiplePages) {
-            int finalWidth = MeasureSpec.getSize(widthMeasureSpec);
             int desiredHeight = 0;
+            int desiredWidth = Math.round(mAvailableWidth);
 
             for (int i = 0; i < getChildCount(); i++) {
                 View child = getChildAt(i);
-                measureChild(child, widthMeasureSpec, heightMeasureSpec);
+                // Measure children based on available height and width.
+                measureChild(child,
+                        MeasureSpec.makeMeasureSpec(desiredWidth, MeasureSpec.EXACTLY),
+                        MeasureSpec.makeMeasureSpec(Math.round(mAvailableHeight),
+                                MeasureSpec.AT_MOST));
                 // Use height of tallest child as we have limited height.
-                desiredHeight = Math.max(desiredHeight, child.getMeasuredHeight());
+                int childHeight = child.getMeasuredHeight();
+                desiredHeight = Math.max(desiredHeight, childHeight);
             }
 
             int finalHeight = resolveSizeAndState(desiredHeight, heightMeasureSpec, 0);
+            int finalWidth = resolveSizeAndState(desiredWidth, widthMeasureSpec, 0);
+
             setMeasuredDimension(finalWidth, finalHeight);
         } else {
             super.onMeasure(widthMeasureSpec, heightMeasureSpec);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 6aaa7d2..3be6faa 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -20,7 +20,6 @@
 import static com.android.launcher3.allapps.ActivityAllAppsContainerView.AdapterHolder.SEARCH;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_SEARCHED;
 import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
-import static com.android.launcher3.widget.picker.WidgetsListAdapter.VIEW_TYPE_WIDGETS_LIST;
 
 import android.animation.Animator;
 import android.content.Context;
@@ -1023,35 +1022,7 @@
                 default:
                     break;
             }
-            mWidgetsListItemAnimator = new DefaultItemAnimator() {
-                @Override
-                public boolean animateChange(RecyclerView.ViewHolder oldHolder,
-                        RecyclerView.ViewHolder newHolder, int fromLeft, int fromTop, int toLeft,
-                        int toTop) {
-                    // As we expand an item, the content / widgets list that appears (with add
-                    // event) also gets change events as its previews load asynchronously. The
-                    // super implementation of animateChange cancels the animations on it - breaking
-                    // the "add animation". Instead, here, we skip "change" animation for content
-                    // list - because we want it to either appear or disappear. And, the previews
-                    // themselves have their own animation when loaded, so, we don't need change
-                    // animations for them anyway. Below, we do-nothing.
-                    if (oldHolder.getItemViewType() == VIEW_TYPE_WIDGETS_LIST) {
-                        dispatchChangeStarting(oldHolder, true);
-                        dispatchChangeFinished(oldHolder, true);
-                        return true;
-                    }
-                    return super.animateChange(oldHolder, newHolder, fromLeft, fromTop, toLeft,
-                            toTop);
-                }
-            };
-            // Disable change animations because it disrupts the item focus upon adapter item
-            // change.
-            mWidgetsListItemAnimator.setSupportsChangeAnimations(false);
-            // Make the moves a bit faster, so that the amount of time for which user sees the
-            // bottom-sheet background before "add" animation starts is less making it smoother.
-            mWidgetsListItemAnimator.setChangeDuration(90);
-            mWidgetsListItemAnimator.setMoveDuration(90);
-            mWidgetsListItemAnimator.setAddDuration(300);
+            mWidgetsListItemAnimator = new WidgetsListItemAnimator();
         }
 
         private int getEmptySpaceHeight() {
@@ -1065,7 +1036,7 @@
             mWidgetsRecyclerView.setClipChildren(false);
             mWidgetsRecyclerView.setAdapter(mWidgetsListAdapter);
             mWidgetsRecyclerView.bindFastScrollbar(mFastScroller);
-            mWidgetsRecyclerView.setItemAnimator(mWidgetsListItemAnimator);
+            mWidgetsRecyclerView.setItemAnimator(isTwoPane() ? null : mWidgetsListItemAnimator);
             mWidgetsRecyclerView.setHeaderViewDimensionsProvider(WidgetsFullSheet.this);
             if (!isTwoPane()) {
                 mWidgetsRecyclerView.setEdgeEffectFactory(
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListItemAnimator.java b/src/com/android/launcher3/widget/picker/WidgetsListItemAnimator.java
new file mode 100644
index 0000000..854700f
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListItemAnimator.java
@@ -0,0 +1,60 @@
+/*
+ * 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.widget.picker;
+
+import static com.android.launcher3.widget.picker.WidgetsListAdapter.VIEW_TYPE_WIDGETS_LIST;
+
+import androidx.recyclerview.widget.DefaultItemAnimator;
+import androidx.recyclerview.widget.RecyclerView;
+
+public class WidgetsListItemAnimator extends DefaultItemAnimator {
+    public static final int CHANGE_DURATION_MS = 90;
+    public static final int MOVE_DURATION_MS = 90;
+    public static final int ADD_DURATION_MS = 120;
+
+    public WidgetsListItemAnimator() {
+        super();
+
+        // Disable change animations because it disrupts the item focus upon adapter item
+        // change.
+        setSupportsChangeAnimations(false);
+        // Make the moves a bit faster, so that the amount of time for which user sees the
+        // bottom-sheet background before "add" animation starts is less making it smoother.
+        setChangeDuration(CHANGE_DURATION_MS);
+        setMoveDuration(MOVE_DURATION_MS);
+        setAddDuration(ADD_DURATION_MS);
+    }
+    @Override
+    public boolean animateChange(RecyclerView.ViewHolder oldHolder,
+            RecyclerView.ViewHolder newHolder, int fromLeft, int fromTop, int toLeft,
+            int toTop) {
+        // As we expand an item, the content / widgets list that appears (with add
+        // event) also gets change events as its previews load asynchronously. The
+        // super implementation of animateChange cancels the animations on it - breaking
+        // the "add animation". Instead, here, we skip "change" animation for content
+        // list - because we want it to either appear or disappear. And, the previews
+        // themselves have their own animation when loaded, so, we don't need change
+        // animations for them anyway. Below, we do-nothing.
+        if (oldHolder.getItemViewType() == VIEW_TYPE_WIDGETS_LIST) {
+            dispatchChangeStarting(oldHolder, true);
+            dispatchChangeFinished(oldHolder, true);
+            return true;
+        }
+        return super.animateChange(oldHolder, newHolder, fromLeft, fromTop, toLeft,
+                toTop);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
index 56352cc..45d733a 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -15,6 +15,11 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.widget.picker.WidgetsListItemAnimator.CHANGE_DURATION_MS;
+import static com.android.launcher3.widget.picker.WidgetsListItemAnimator.MOVE_DURATION_MS;
+
+import static android.animation.ValueAnimator.areAnimatorsEnabled;
+
 import android.content.Context;
 import android.graphics.Bitmap;
 import android.util.Log;
@@ -26,7 +31,6 @@
 import android.view.View.OnLongClickListener;
 import android.view.ViewGroup;
 import android.widget.TableLayout;
-import android.widget.TableRow;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Px;
@@ -36,6 +40,7 @@
 import com.android.launcher3.recyclerview.ViewHolderBinder;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.WidgetTableRow;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.util.WidgetsTableUtils;
 
@@ -111,17 +116,15 @@
         for (int i = 0; i < widgetItemsTable.size(); i++) {
             List<WidgetItem> widgetItemsPerRow = widgetItemsTable.get(i);
             for (int j = 0; j < widgetItemsPerRow.size(); j++) {
-                TableRow row = (TableRow) table.getChildAt(i);
+                WidgetTableRow row = (WidgetTableRow) table.getChildAt(i);
                 row.setVisibility(View.VISIBLE);
                 WidgetCell widget = (WidgetCell) row.getChildAt(j);
                 widget.clear();
+                widget.addPreviewReadyListener(row);
                 WidgetItem widgetItem = widgetItemsPerRow.get(j);
                 widget.setVisibility(View.VISIBLE);
 
-                // When preview loads, notify adapter to rebind the item and possibly animate
-                widget.applyFromCellItem(widgetItem,
-                        bitmap -> holder.onPreviewLoaded(Pair.create(widgetItem, bitmap)),
-                        holder.previewCache.get(widgetItem));
+                widget.applyFromCellItem(widgetItem);
                 widget.requestLayout();
             }
         }
@@ -143,14 +146,19 @@
 
         for (int i = 0; i < widgetItemsTable.size(); i++) {
             List<WidgetItem> widgetItems = widgetItemsTable.get(i);
-            TableRow tableRow;
+            WidgetTableRow tableRow;
             if (i < table.getChildCount()) {
-                tableRow = (TableRow) table.getChildAt(i);
+                tableRow = (WidgetTableRow) table.getChildAt(i);
             } else {
-                tableRow = new TableRow(table.getContext());
+                tableRow = new WidgetTableRow(table.getContext());
                 tableRow.setGravity(Gravity.TOP);
                 table.addView(tableRow);
             }
+            // Pass resize delay to let the "move" and "change" animations run before resizing the
+            // row.
+            tableRow.setupRow(widgetItems.size(),
+                    /*resizeDelayMs=*/
+                    areAnimatorsEnabled() ? (CHANGE_DURATION_MS + MOVE_DURATION_MS) : 0);
             if (tableRow.getChildCount() > widgetItems.size()) {
                 for (int j = widgetItems.size(); j < tableRow.getChildCount(); j++) {
                     tableRow.getChildAt(j).setVisibility(View.GONE);
@@ -161,6 +169,7 @@
                             R.layout.widget_cell, tableRow, false);
                     // set up touch.
                     widget.setOnClickListener(mIconClickListener);
+                    widget.addPreviewReadyListener(tableRow);
                     View preview = widget.findViewById(R.id.widget_preview_container);
                     preview.setOnClickListener(mIconClickListener);
                     preview.setOnLongClickListener(mIconLongClickListener);
@@ -176,7 +185,7 @@
         int numOfRows = holder.tableContainer.getChildCount();
         holder.previewCache.clear();
         for (int i = 0; i < numOfRows; i++) {
-            TableRow tableRow = (TableRow) holder.tableContainer.getChildAt(i);
+            WidgetTableRow tableRow = (WidgetTableRow) holder.tableContainer.getChildAt(i);
             int numOfCols = tableRow.getChildCount();
             for (int j = 0; j < numOfCols; j++) {
                 WidgetCell widget = (WidgetCell) tableRow.getChildAt(j);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
index 6dbad5c..1ed3d88 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
@@ -27,9 +27,7 @@
 import android.view.Gravity;
 import android.view.LayoutInflater;
 import android.view.View;
-import android.view.ViewGroup;
 import android.widget.TableLayout;
-import android.widget.TableRow;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.Px;
@@ -38,6 +36,7 @@
 import com.android.launcher3.R;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.WidgetTableRow;
 import com.android.launcher3.widget.picker.util.WidgetPreviewContainerSize;
 
 import java.util.ArrayList;
@@ -105,7 +104,8 @@
 
         for (int i = 0; i < recommendationTable.size(); i++) {
             List<WidgetItem> widgetItems = recommendationTable.get(i);
-            TableRow tableRow = new TableRow(getContext());
+            WidgetTableRow tableRow = new WidgetTableRow(getContext());
+            tableRow.setupRow(widgetItems.size(), /*resizeDelayMs=*/ 0);
             tableRow.setGravity(Gravity.TOP);
             for (WidgetItem widgetItem : widgetItems) {
                 WidgetCell widgetCell = addItemCell(tableRow);
@@ -121,9 +121,10 @@
         setVisibility(VISIBLE);
     }
 
-    private WidgetCell addItemCell(ViewGroup parent) {
+    private WidgetCell addItemCell(WidgetTableRow parent) {
         WidgetCell widget = (WidgetCell) LayoutInflater.from(
                 getContext()).inflate(R.layout.widget_cell, parent, false);
+        widget.addPreviewReadyListener(parent);
         widget.setOnClickListener(mWidgetCellOnClickListener);
 
         View previewContainer = widget.findViewById(R.id.widget_preview_container);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index f835e18..5d71db6 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -67,7 +67,7 @@
 
     // This ratio defines the max percentage of content area that the recommendations can display
     // with respect to the bottom sheet's height.
-    private static final float RECOMMENDATION_SECTION_HEIGHT_RATIO_TWO_PANE = 0.60f;
+    private static final float RECOMMENDATION_SECTION_HEIGHT_RATIO_TWO_PANE = 0.70f;
     private FrameLayout mSuggestedWidgetsContainer;
     private WidgetsListHeader mSuggestedWidgetsHeader;
     private PackageUserKey mSuggestedWidgetsPackageUserKey;
diff --git a/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java b/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java
index 5c50e97..6ce1cb7 100644
--- a/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java
+++ b/tests/src/com/android/launcher3/allapps/PrivateProfileManagerTest.java
@@ -208,7 +208,7 @@
         ArgumentCaptor<Intent> acIntent = ArgumentCaptor.forClass(Intent.class);
         mPrivateProfileManager.setPrivateSpaceSettingsAvailable(true);
 
-        mPrivateProfileManager.openPrivateSpaceSettings(null);
+        mContext.startActivity(expectedIntent);
 
         Mockito.verify(mContext).startActivity(acIntent.capture());
         assertEquals("Intent Action is different",
diff --git a/tests/src/com/android/launcher3/allapps/PrivateSpaceSettingsButtonTest.java b/tests/src/com/android/launcher3/allapps/PrivateSpaceSettingsButtonTest.java
new file mode 100644
index 0000000..9537e1c
--- /dev/null
+++ b/tests/src/com/android/launcher3/allapps/PrivateSpaceSettingsButtonTest.java
@@ -0,0 +1,56 @@
+/*
+ * 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.allapps;
+
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PRIVATESPACE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.util.ActivityContextWrapper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(AndroidJUnit4.class)
+public class PrivateSpaceSettingsButtonTest {
+
+    private PrivateSpaceSettingsButton mVut;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        Context context = new ActivityContextWrapper(getApplicationContext());
+        mVut = new PrivateSpaceSettingsButton(context);
+    }
+
+    @Test
+    public void privateSpaceSettingsAppInfo_hasCorrectIdAndContainer() {
+        AppInfo appInfo = mVut.createPrivateSpaceSettingsAppInfo();
+
+        assertThat(appInfo.id).isEqualTo(CONTAINER_PRIVATESPACE);
+        assertThat(appInfo.container).isEqualTo(CONTAINER_PRIVATESPACE);
+    }
+}