Merge "Cleanup task view when bubble is removed" into main
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
index f32f030..50a58da 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
@@ -465,7 +465,7 @@
     /**
      * Call when all the views should be removed/cleaned up.
      */
-    void cleanupViews() {
+    public void cleanupViews() {
         cleanupViews(true);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index aea3ca1..d0db708 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -730,7 +730,7 @@
             // window to show this in, but we use a separate code path.
             // TODO(b/273312602): consider foldables where we do need a stack view when folded
             if (mLayerView == null) {
-                mLayerView = new BubbleBarLayerView(mContext, this);
+                mLayerView = new BubbleBarLayerView(mContext, this, mBubbleData);
                 mLayerView.setUnBubbleConversationCallback(mSysuiProxy::onUnbubbleConversation);
             }
         } else {
@@ -1714,8 +1714,7 @@
         @Override
         public void removeBubble(Bubble removedBubble) {
             if (mLayerView != null) {
-                // TODO: need to check if there's something that needs to happen here, e.g. if
-                //  the currently selected & expanded bubble is removed?
+                mLayerView.removeBubble(removedBubble);
             }
         }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
index 127c7e8..dbfa260 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
@@ -15,7 +15,6 @@
  */
 package com.android.wm.shell.bubbles;
 
-import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
 import static com.android.wm.shell.bubbles.Bubble.KEY_APP_BUBBLE;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.DEBUG_BUBBLE_DATA;
@@ -41,6 +40,7 @@
 import com.android.launcher3.icons.BubbleIconFactory;
 import com.android.wm.shell.R;
 import com.android.wm.shell.bubbles.Bubbles.DismissReason;
+import com.android.wm.shell.bubbles.bar.BubbleBarLayerView;
 import com.android.wm.shell.common.bubbles.BubbleBarUpdate;
 import com.android.wm.shell.common.bubbles.RemovedBubble;
 
@@ -427,7 +427,7 @@
     /**
      * When this method is called it is expected that all info in the bubble has completed loading.
      * @see Bubble#inflate(BubbleViewInfoTask.Callback, Context, BubbleController, BubbleStackView,
-     * BubbleIconFactory, boolean)
+     * BubbleBarLayerView, BubbleIconFactory, boolean)
      */
     void notificationEntryUpdated(Bubble bubble, boolean suppressFlyout, boolean showInShade) {
         if (DEBUG_BUBBLE_DATA) {
@@ -1069,7 +1069,6 @@
     /**
      * The set of bubbles in row.
      */
-    @VisibleForTesting(visibility = PACKAGE)
     public List<Bubble> getBubbles() {
         return Collections.unmodifiableList(mBubbles);
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskViewHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskViewHelper.java
index 5fc67d7..dc27133 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskViewHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskViewHelper.java
@@ -31,6 +31,7 @@
 import android.graphics.Rect;
 import android.util.Log;
 import android.view.View;
+import android.view.ViewGroup;
 
 import androidx.annotation.Nullable;
 
@@ -186,6 +187,7 @@
             }
             if (mTaskView != null) {
                 mTaskView.release();
+                ((ViewGroup) mParentView).removeView(mTaskView);
                 mTaskView = null;
             }
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
index 00d683e..73a9cf4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
@@ -266,13 +266,8 @@
         mListener.onBackPressed();
     }
 
-    /** Cleans up task view, should be called when the bubble is no longer active. */
+    /** Cleans up the expanded view, should be called when the bubble is no longer active. */
     public void cleanUpExpandedState() {
-        if (mBubbleTaskViewHelper != null) {
-            if (mTaskView != null) {
-                removeView(mTaskView);
-            }
-        }
         mMenuViewController.hideMenu(false /* animated */);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
index bd8ce80..b95d258 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
@@ -34,7 +34,9 @@
 import android.widget.FrameLayout;
 
 import com.android.wm.shell.R;
+import com.android.wm.shell.bubbles.Bubble;
 import com.android.wm.shell.bubbles.BubbleController;
+import com.android.wm.shell.bubbles.BubbleData;
 import com.android.wm.shell.bubbles.BubbleOverflow;
 import com.android.wm.shell.bubbles.BubblePositioner;
 import com.android.wm.shell.bubbles.BubbleViewProvider;
@@ -61,6 +63,7 @@
     private static final float SCRIM_ALPHA = 0.2f;
 
     private final BubbleController mBubbleController;
+    private final BubbleData mBubbleData;
     private final BubblePositioner mPositioner;
     private final BubbleBarAnimationHelper mAnimationHelper;
     private final BubbleEducationViewController mEducationViewController;
@@ -85,9 +88,10 @@
     private TouchDelegate mHandleTouchDelegate;
     private final Rect mHandleTouchBounds = new Rect();
 
-    public BubbleBarLayerView(Context context, BubbleController controller) {
+    public BubbleBarLayerView(Context context, BubbleController controller, BubbleData bubbleData) {
         super(context);
         mBubbleController = controller;
+        mBubbleData = bubbleData;
         mPositioner = mBubbleController.getPositioner();
 
         mAnimationHelper = new BubbleBarAnimationHelper(context,
@@ -236,15 +240,44 @@
         showScrim(true);
     }
 
+    /** Removes the given {@code bubble}. */
+    public void removeBubble(Bubble bubble) {
+        if (mBubbleData.getBubbles().isEmpty()) {
+            // we're removing the last bubble. collapse the expanded view and cleanup bubble views
+            // at the end.
+            collapse(bubble::cleanupViews);
+        } else {
+            bubble.cleanupViews();
+        }
+    }
+
     /** Collapses any showing expanded view */
     public void collapse() {
+        collapse(/* endAction= */ null);
+    }
+
+    /**
+     * Collapses any showing expanded view.
+     *
+     * @param endAction an action to run and the end of the collapse animation.
+     */
+    public void collapse(@Nullable Runnable endAction) {
+        if (!mIsExpanded) {
+            return;
+        }
         mIsExpanded = false;
         final BubbleBarExpandedView viewToRemove = mExpandedView;
         mEducationViewController.hideEducation(/* animated = */ true);
+        Runnable runnable = () -> {
+            removeView(viewToRemove);
+            if (endAction != null) {
+                endAction.run();
+            }
+        };
         if (mDragController != null && mDragController.isStuckToDismiss()) {
-            mAnimationHelper.animateDismiss(() -> removeView(viewToRemove));
+            mAnimationHelper.animateDismiss(runnable);
         } else {
-            mAnimationHelper.animateCollapse(() -> removeView(viewToRemove));
+            mAnimationHelper.animateCollapse(runnable);
         }
         mBubbleController.getSysuiProxy().onStackExpandChanged(false);
         mExpandedView = null;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleViewInfoTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleViewInfoTest.kt
index 75965d6..1668e37 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleViewInfoTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleViewInfoTest.kt
@@ -143,7 +143,7 @@
                 bubbleController,
                 mainExecutor
             )
-        bubbleBarLayerView = BubbleBarLayerView(context, bubbleController)
+        bubbleBarLayerView = BubbleBarLayerView(context, bubbleController, bubbleData)
     }
 
     @Test