Merge "Revert "Separate desktop and fullscreen carousel"" into main
diff --git a/quickstep/Android.bp b/quickstep/Android.bp
index 4c724dc..2ef9f82 100644
--- a/quickstep/Android.bp
+++ b/quickstep/Android.bp
@@ -64,3 +64,11 @@
         "tests/multivalentScreenshotTests/src/**/*.kt",
     ],
 }
+
+filegroup {
+    name: "launcher3-quickstep-testing",
+    path: "testing",
+    srcs: [
+        "testing/**/*.kt",
+    ],
+}
diff --git a/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java b/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java
index e940553..a63ba0f 100644
--- a/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java
+++ b/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java
@@ -26,7 +26,6 @@
 import android.animation.AnimatorSet;
 import android.content.Context;
 import android.os.Handler;
-import android.os.RemoteException;
 import android.util.Log;
 import android.view.IRemoteAnimationFinishedCallback;
 import android.view.RemoteAnimationTarget;
@@ -210,7 +209,7 @@
          * animation finished runnable.
          */
         @Override
-        public void onAnimationFinished() throws RemoteException {
+        public void onAnimationFinished() {
             mASyncFinishRunnable.run();
         }
     }
@@ -240,12 +239,5 @@
         @Override
         @UiThread
         default void onAnimationCancelled() {}
-
-        /**
-         * Returns whether this animation factory supports a tightly coupled return animation.
-         */
-        default boolean supportsReturnTransition() {
-            return false;
-        }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 2457cfd..da6c05e 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -111,6 +111,7 @@
 import android.view.animation.PathInterpolator;
 import android.window.RemoteTransition;
 import android.window.TransitionFilter;
+import android.window.WindowAnimationState;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -141,6 +142,9 @@
 import com.android.quickstep.RemoteAnimationTargets;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TaskViewUtils;
+import com.android.quickstep.util.AlreadyStartedBackAnimState;
+import com.android.quickstep.util.AnimatorBackState;
+import com.android.quickstep.util.BackAnimState;
 import com.android.quickstep.util.MultiValueUpdateListener;
 import com.android.quickstep.util.RectFSpringAnim;
 import com.android.quickstep.util.RectFSpringAnim.DefaultSpringConfig;
@@ -176,9 +180,6 @@
  */
 public class QuickstepTransitionManager implements OnDeviceProfileChangeListener {
 
-    private static final String TRANSITION_COOKIE_PREFIX =
-            "com.android.launcher3.QuickstepTransitionManager_activityLaunch";
-
     private static final boolean ENABLE_SHELL_STARTING_SURFACE =
             SystemProperties.getBoolean("persist.debug.shell_starting_surface", true);
 
@@ -347,14 +348,7 @@
         IRemoteCallback endCallback = completeRunnableListCallback(onEndCallback);
         options.setOnAnimationAbortListener(endCallback);
         options.setOnAnimationFinishedListener(endCallback);
-
-        IBinder cookie = mAppLaunchRunner.supportsReturnTransition()
-                ? ((ContainerAnimationRunner) mAppLaunchRunner).getCookie() : null;
-        addLaunchCookie(cookie, itemInfo, options);
-
-        // Register the return animation so it can be triggered on back from the app to home.
-        maybeRegisterAppReturnTransition(v);
-
+        options.setLaunchCookie(StableViewInfo.toLaunchCookie(itemInfo));
         return new ActivityOptionsWrapper(options, onEndCallback);
     }
 
@@ -367,21 +361,9 @@
         ItemInfo tag = (ItemInfo) v.getTag();
         ContainerAnimationRunner containerRunner = null;
         if (tag != null && tag.shouldUseBackgroundAnimation()) {
-            // The cookie should only override the default used by launcher if container return
-            // animations are enabled.
-            ActivityTransitionAnimator.TransitionCookie cookie =
-                    checkReturnAnimationsFlags()
-                            ? new ActivityTransitionAnimator.TransitionCookie(
-                                    TRANSITION_COOKIE_PREFIX + tag.id)
-                            : null;
-            ContainerAnimationRunner launchAnimationRunner =
-                    ContainerAnimationRunner.fromView(
-                            v, cookie, true /* forLaunch */, mLauncher, mStartingWindowListener,
-                            onEndCallback);
-
-            if (launchAnimationRunner != null) {
-                containerRunner = launchAnimationRunner;
-            }
+            containerRunner = ContainerAnimationRunner.fromView(
+                    v, true /* forLaunch */, mLauncher, mStartingWindowListener, onEndCallback,
+                    null /* windowState */);
         }
 
         mAppLaunchRunner = containerRunner != null
@@ -391,51 +373,6 @@
     }
 
     /**
-     * If container return animations are enabled and the current launch runner is itself a
-     * {@link ContainerAnimationRunner}, registers a matching return animation that de-registers
-     * itself after it has run once or is made obsolete by the view going away.
-     */
-    private void maybeRegisterAppReturnTransition(View v) {
-        if (!checkReturnAnimationsFlags() || !mAppLaunchRunner.supportsReturnTransition()) {
-            return;
-        }
-
-        ActivityTransitionAnimator.TransitionCookie cookie =
-                ((ContainerAnimationRunner) mAppLaunchRunner).getCookie();
-        RunnableList onEndCallback = new RunnableList();
-        ContainerAnimationRunner runner =
-                ContainerAnimationRunner.fromView(
-                        v, cookie, false /* forLaunch */, mLauncher, mStartingWindowListener,
-                        onEndCallback);
-        RemoteTransition transition =
-                new RemoteTransition(
-                        new LauncherAnimationRunner(
-                                mHandler, runner, true /* startAtFrontOfQueue */
-                        ).toRemoteTransition()
-                );
-
-        SystemUiProxy.INSTANCE.get(mLauncher).registerRemoteTransition(
-                transition, ContainerAnimationRunner.buildBackToHomeFilter(cookie, mLauncher));
-        ContainerAnimationRunner.setUpRemoteAnimationCleanup(
-                v, transition, onEndCallback, mLauncher);
-    }
-
-    /**
-     * Adds a new launch cookie for the activity launch if supported.
-     * Prioritizes the explicitly provided cookie, falling back on extracting one from the given
-     * {@link ItemInfo} if necessary.
-     */
-    private void addLaunchCookie(IBinder cookie, ItemInfo info, ActivityOptions options) {
-        if (cookie == null) {
-            cookie = StableViewInfo.toLaunchCookie(info);
-        }
-
-        if (cookie != null) {
-            options.setLaunchCookie(cookie);
-        }
-    }
-
-    /**
      * Whether the launch is a recents app transition and we should do a launch animation
      * from the recents view. Note that if the remote animation targets are not provided, this
      * may not always be correct as we may resolve the opening app to a task when the animation
@@ -1613,19 +1550,48 @@
      * Creates the {@link RectFSpringAnim} and {@link AnimatorSet} required to animate
      * the transition.
      */
-    public Pair<RectFSpringAnim, AnimatorSet> createWallpaperOpenAnimations(
+    @NonNull
+    public BackAnimState createWallpaperOpenAnimations(
             RemoteAnimationTarget[] appTargets,
-            RemoteAnimationTarget[] wallpaperTargets,
+            RemoteAnimationTarget[] wallpapers,
+            RemoteAnimationTarget[] nonAppTargets,
             RectF startRect,
             float startWindowCornerRadius,
             boolean fromPredictiveBack) {
+        View launcherView = findLauncherView(appTargets);
+        if (checkReturnAnimationsFlags()
+                && launcherView != null
+                && launcherView.getTag() instanceof ItemInfo info
+                && info.shouldUseBackgroundAnimation()) {
+            // Try to create a return animation
+            RunnableList onEndCallback = new RunnableList();
+            WindowAnimationState windowState = new WindowAnimationState();
+            windowState.bounds = startRect;
+            windowState.bottomLeftRadius = windowState.bottomRightRadius =
+                    windowState.topLeftRadius = windowState.topRightRadius =
+                            startWindowCornerRadius;
+            ContainerAnimationRunner runner = ContainerAnimationRunner.fromView(
+                    launcherView, false /* forLaunch */, mLauncher, mStartingWindowListener,
+                    onEndCallback, windowState);
+            if (runner != null) {
+                runner.startAnimation(TRANSIT_CLOSE,
+                        appTargets, wallpapers, nonAppTargets,
+                        new IRemoteAnimationFinishedCallback.Stub() {
+                            @Override
+                            public void onAnimationFinished() {
+                                onEndCallback.executeAllAndDestroy();
+                            }
+                        });
+                return new AlreadyStartedBackAnimState(onEndCallback);
+            }
+        }
+
         AnimatorSet anim = new AnimatorSet();
         RectFSpringAnim rectFSpringAnim = null;
 
         final boolean launcherIsForceInvisibleOrOpening = mLauncher.isForceInvisible()
                 || launcherIsATargetWithMode(appTargets, MODE_OPENING);
 
-        View launcherView = findLauncherView(appTargets);
         boolean playFallBackAnimation = (launcherView == null
                 && launcherIsForceInvisibleOrOpening)
                 || mLauncher.getWorkspace().isOverlayShown()
@@ -1725,7 +1691,7 @@
             }
         }
 
-        return new Pair(rectFSpringAnim, anim);
+        return new AnimatorBackState(rectFSpringAnim, anim);
     }
 
     public static int getTaskbarToHomeDuration() {
@@ -1775,14 +1741,14 @@
                 }
             }
 
-            Pair<RectFSpringAnim, AnimatorSet> pair = createWallpaperOpenAnimations(
-                    appTargets, wallpaperTargets, resolveRectF,
+            BackAnimState bankAnimState = createWallpaperOpenAnimations(
+                    appTargets, wallpaperTargets, nonAppTargets, resolveRectF,
                     QuickStepContract.getWindowCornerRadius(mLauncher),
                     false /* fromPredictiveBack */);
 
             TaskViewUtils.createSplitAuxiliarySurfacesAnimator(nonAppTargets, false, null);
             mLauncher.clearForceInvisibleFlag(INVISIBLE_ALL);
-            result.setAnimation(pair.second, mLauncher);
+            bankAnimState.applyToAnimationResult(result, mLauncher);
         }
     }
 
@@ -1850,29 +1816,19 @@
         /** The delegate runner that handles the actual animation. */
         private final RemoteAnimationDelegate<IRemoteAnimationFinishedCallback> mDelegate;
 
-        @Nullable
-        private final ActivityTransitionAnimator.TransitionCookie mCookie;
-
         private ContainerAnimationRunner(
-                RemoteAnimationDelegate<IRemoteAnimationFinishedCallback> delegate,
-                ActivityTransitionAnimator.TransitionCookie cookie) {
+                RemoteAnimationDelegate<IRemoteAnimationFinishedCallback> delegate) {
             mDelegate = delegate;
-            mCookie = cookie;
-        }
-
-        @Nullable
-        ActivityTransitionAnimator.TransitionCookie getCookie() {
-            return mCookie;
         }
 
         @Nullable
         static ContainerAnimationRunner fromView(
                 View v,
-                ActivityTransitionAnimator.TransitionCookie cookie,
                 boolean forLaunch,
                 Launcher launcher,
                 StartingWindowListener startingWindowListener,
-                RunnableList onEndCallback) {
+                RunnableList onEndCallback,
+                @Nullable WindowAnimationState windowState) {
             if (!forLaunch && !checkReturnAnimationsFlags()) {
                 throw new IllegalStateException(
                         "forLaunch cannot be false when the enableContainerReturnAnimations or "
@@ -1882,7 +1838,7 @@
             // First the controller is created. This is used by the runner to animate the
             // origin/target view.
             ActivityTransitionAnimator.Controller controller =
-                    buildController(v, cookie, forLaunch);
+                    buildController(v, forLaunch, windowState);
             if (controller == null) {
                 return null;
             }
@@ -1907,8 +1863,7 @@
 
             return new ContainerAnimationRunner(
                     new ActivityTransitionAnimator.AnimationDelegate(
-                            MAIN_EXECUTOR, controller, callback, listener),
-                    cookie);
+                            MAIN_EXECUTOR, controller, callback, listener));
         }
 
         /**
@@ -1918,7 +1873,7 @@
          */
         @Nullable
         private static ActivityTransitionAnimator.Controller buildController(
-                View v, ActivityTransitionAnimator.TransitionCookie cookie, boolean isLaunching) {
+                View v, boolean isLaunching, @Nullable WindowAnimationState windowState) {
             View viewToUse = findLaunchableViewWithBackground(v);
             if (viewToUse == null) {
                 return null;
@@ -1949,8 +1904,8 @@
 
                 @Nullable
                 @Override
-                public ActivityTransitionAnimator.TransitionCookie getTransitionCookie() {
-                    return cookie;
+                public WindowAnimationState getWindowAnimatorState() {
+                    return windowState;
                 }
             };
         }
@@ -1964,81 +1919,26 @@
                 View view) {
             View current = view;
             while (current.getBackground() == null || !(current instanceof LaunchableView)) {
-                if (!(current.getParent() instanceof View)) {
+                if (current.getParent() instanceof View v) {
+                    current = v;
+                } else {
                     return null;
                 }
-
-                current = (View) current.getParent();
             }
-
             return (T) current;
         }
 
-        /**
-         * Builds the filter used by WM Shell to match app closing transitions (only back, no home
-         * button/gesture) to the given launch cookie.
-         */
-        static TransitionFilter buildBackToHomeFilter(
-                ActivityTransitionAnimator.TransitionCookie cookie, Launcher launcher) {
-            // Closing activity must include the cookie in its list of launch cookies.
-            TransitionFilter.Requirement appRequirement = new TransitionFilter.Requirement();
-            appRequirement.mActivityType = ACTIVITY_TYPE_STANDARD;
-            appRequirement.mLaunchCookie = cookie;
-            appRequirement.mModes = new int[]{TRANSIT_CLOSE, TRANSIT_TO_BACK};
-            // Opening activity must be Launcher.
-            TransitionFilter.Requirement launcherRequirement = new TransitionFilter.Requirement();
-            launcherRequirement.mActivityType = ACTIVITY_TYPE_HOME;
-            launcherRequirement.mModes = new int[]{TRANSIT_OPEN, TRANSIT_TO_FRONT};
-            launcherRequirement.mTopActivity = launcher.getComponentName();
-            // Transition types CLOSE and TO_BACK match the back button/gesture but not the  home
-            // button/gesture.
-            TransitionFilter filter = new TransitionFilter();
-            filter.mTypeSet = new int[]{TRANSIT_CLOSE, TRANSIT_TO_BACK};
-            filter.mRequirements =
-                    new TransitionFilter.Requirement[]{appRequirement, launcherRequirement};
-            return filter;
-        }
-
-        /**
-         * Creates various conditions to ensure that the given transition is cleaned up correctly
-         * when necessary:
-         * - if the transition has run, it is the callback that unregisters it;
-         * - if the associated view is detached before the transition has had an opportunity to run,
-         *   a {@link View.OnAttachStateChangeListener} allows us to do the same (and removes
-         *   itself).
-         */
-        static void setUpRemoteAnimationCleanup(
-                View v, RemoteTransition transition, RunnableList callback, Launcher launcher) {
-            View.OnAttachStateChangeListener listener = new View.OnAttachStateChangeListener() {
-                @Override
-                public void onViewAttachedToWindow(@NonNull View v) {}
-
-                @Override
-                public void onViewDetachedFromWindow(@NonNull View v) {
-                    SystemUiProxy.INSTANCE.get(launcher)
-                            .unregisterRemoteTransition(transition);
-                    v.removeOnAttachStateChangeListener(this);
-                }
-            };
-
-            // Remove the animation as soon as it has run once.
-            callback.add(() -> {
-                SystemUiProxy.INSTANCE.get(launcher).unregisterRemoteTransition(transition);
-                if (v != null) {
-                    v.removeOnAttachStateChangeListener(listener);
-                }
-            });
-
-            // Remove the animation when the view is detached from the hierarchy.
-            // This is so that if back is not invoked (e.g. if we go back home through the home
-            // gesture) we don't have obsolete transitions staying registered.
-            v.addOnAttachStateChangeListener(listener);
-        }
-
         @Override
         public void onAnimationStart(int transit, RemoteAnimationTarget[] appTargets,
                 RemoteAnimationTarget[] wallpaperTargets, RemoteAnimationTarget[] nonAppTargets,
                 LauncherAnimationRunner.AnimationResult result) {
+            startAnimation(
+                    transit, appTargets, wallpaperTargets, nonAppTargets, result);
+        }
+
+        public void startAnimation(int transit, RemoteAnimationTarget[] appTargets,
+                RemoteAnimationTarget[] wallpaperTargets, RemoteAnimationTarget[] nonAppTargets,
+                IRemoteAnimationFinishedCallback result) {
             mDelegate.onAnimationStart(
                     transit, appTargets, wallpaperTargets, nonAppTargets, result);
         }
@@ -2047,11 +1947,6 @@
         public void onAnimationCancelled() {
             mDelegate.onAnimationCancelled();
         }
-
-        @Override
-        public boolean supportsReturnTransition() {
-            return true;
-        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
index 360c216..1124aac 100644
--- a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
+++ b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
@@ -27,7 +27,6 @@
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
-import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
 import android.content.ComponentCallbacks;
 import android.content.res.Configuration;
@@ -38,7 +37,6 @@
 import android.os.Handler;
 import android.os.RemoteException;
 import android.util.Log;
-import android.util.Pair;
 import android.view.Choreographer;
 import android.view.IRemoteAnimationFinishedCallback;
 import android.view.IRemoteAnimationRunner;
@@ -63,7 +61,7 @@
 import com.android.launcher3.taskbar.LauncherTaskbarUIController;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
-import com.android.quickstep.util.RectFSpringAnim;
+import com.android.quickstep.util.BackAnimState;
 import com.android.systemui.shared.system.QuickStepContract;
 
 import java.lang.ref.WeakReference;
@@ -109,8 +107,6 @@
     private RemoteAnimationTarget mLauncherTarget;
     private View mLauncherTargetView;
     private final SurfaceControl.Transaction mTransaction = new SurfaceControl.Transaction();
-    private boolean mSpringAnimationInProgress = false;
-    private boolean mAnimatorSetInProgress = false;
     private float mBackProgress = 0;
     private boolean mBackInProgress = false;
     private OnBackInvokedCallbackStub mBackCallback;
@@ -448,14 +444,15 @@
         mQuickstepTransitionManager.transferRectToTargetCoordinate(
                 mBackTarget, mCurrentRect, true, resolveRectF);
 
-        Pair<RectFSpringAnim, AnimatorSet> pair =
+        BackAnimState backAnim =
                 mQuickstepTransitionManager.createWallpaperOpenAnimations(
                     new RemoteAnimationTarget[]{mBackTarget},
                     new RemoteAnimationTarget[0],
+                    new RemoteAnimationTarget[0],
                     resolveRectF,
                     cornerRadius,
                     mBackInProgress /* fromPredictiveBack */);
-        startTransitionAnimations(pair.first, pair.second);
+        startTransitionAnimations(backAnim);
         mLauncher.clearForceInvisibleFlag(INVISIBLE_ALL);
         customizeStatusBarAppearance(true);
     }
@@ -470,8 +467,6 @@
         mCurrentRect.setEmpty();
         mStartRect.setEmpty();
         mInitialTouchPos.set(0, 0);
-        mAnimatorSetInProgress = false;
-        mSpringAnimationInProgress = false;
         setLauncherTargetViewVisible(true);
         mLauncherTargetView = null;
         // We don't call customizeStatusBarAppearance here to prevent the status bar update with
@@ -494,27 +489,8 @@
         }
     }
 
-    private void startTransitionAnimations(RectFSpringAnim springAnim, AnimatorSet anim) {
-        mAnimatorSetInProgress = anim != null;
-        mSpringAnimationInProgress = springAnim != null;
-        if (springAnim != null) {
-            springAnim.addAnimatorListener(
-                    new AnimatorListenerAdapter() {
-                        @Override
-                        public void onAnimationEnd(Animator animation) {
-                            mSpringAnimationInProgress = false;
-                            tryFinishBackAnimation();
-                        }
-                    }
-            );
-        }
-        anim.addListener(new AnimatorListenerAdapter() {
-            @Override
-            public void onAnimationEnd(Animator animation) {
-                mAnimatorSetInProgress = false;
-                tryFinishBackAnimation();
-            }
-        });
+    private void startTransitionAnimations(BackAnimState backAnim) {
+        backAnim.addOnAnimCompleteCallback(this::finishAnimation);
         if (mScrimLayer == null) {
             // Scrim hasn't been attached yet. Let's attach it.
             addScrimLayer();
@@ -534,7 +510,7 @@
             }
         });
         mScrimAlphaAnimator.setDuration(SCRIM_FADE_DURATION).start();
-        anim.start();
+        backAnim.start();
     }
 
     private void loadResources() {
@@ -567,12 +543,6 @@
         mScrimAlpha = 0;
     }
 
-    private void tryFinishBackAnimation() {
-        if (!mSpringAnimationInProgress && !mAnimatorSetInProgress) {
-            finishAnimation();
-        }
-    }
-
     private void customizeStatusBarAppearance(boolean overridingStatusBarFlags) {
         if (mOverridingStatusBarFlags == overridingStatusBarFlags) {
             return;
diff --git a/quickstep/src/com/android/quickstep/util/BackAnimState.kt b/quickstep/src/com/android/quickstep/util/BackAnimState.kt
new file mode 100644
index 0000000..9009eaa
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/BackAnimState.kt
@@ -0,0 +1,70 @@
+/*
+ * 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.quickstep.util
+
+import android.animation.AnimatorSet
+import android.content.Context
+import com.android.launcher3.LauncherAnimationRunner.AnimationResult
+import com.android.launcher3.anim.AnimatorListeners.forEndCallback
+import com.android.launcher3.util.RunnableList
+
+/** Interface to represent animation for back to Launcher transition */
+interface BackAnimState {
+
+    fun addOnAnimCompleteCallback(r: Runnable)
+
+    fun applyToAnimationResult(result: AnimationResult, c: Context)
+
+    fun start()
+}
+
+class AnimatorBackState(private val springAnim: RectFSpringAnim?, private val anim: AnimatorSet?) :
+    BackAnimState {
+
+    override fun addOnAnimCompleteCallback(r: Runnable) {
+        val springAnimWait = RunnableList()
+        springAnim?.addAnimatorListener(forEndCallback(springAnimWait::executeAllAndDestroy))
+            ?: springAnimWait.executeAllAndDestroy()
+
+        val animWait = RunnableList()
+        anim?.addListener(
+            forEndCallback(Runnable { springAnimWait.add(animWait::executeAllAndDestroy) })
+        ) ?: springAnimWait.add(animWait::executeAllAndDestroy)
+        animWait.add(r)
+    }
+
+    override fun applyToAnimationResult(result: AnimationResult, c: Context) {
+        result.setAnimation(anim, c)
+    }
+
+    override fun start() {
+        anim?.start()
+    }
+}
+
+class AlreadyStartedBackAnimState(private val onEndCallback: RunnableList) : BackAnimState {
+
+    override fun addOnAnimCompleteCallback(r: Runnable) {
+        onEndCallback.add(r)
+    }
+
+    override fun applyToAnimationResult(result: AnimationResult, c: Context) {
+        addOnAnimCompleteCallback(result::onAnimationFinished)
+    }
+
+    override fun start() {}
+}
diff --git a/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt b/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
new file mode 100644
index 0000000..37a07c3
--- /dev/null
+++ b/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
@@ -0,0 +1,78 @@
+/*
+ * 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.taskbar.bubbles.testing
+
+import android.app.Notification
+import android.content.Context
+import android.graphics.Bitmap
+import android.graphics.Canvas
+import android.graphics.Color
+import android.graphics.Paint
+import android.util.PathParser
+import android.view.LayoutInflater
+import android.view.ViewGroup
+import com.android.launcher3.R
+import com.android.launcher3.taskbar.bubbles.BubbleBarBubble
+import com.android.launcher3.taskbar.bubbles.BubbleView
+import com.android.wm.shell.shared.bubbles.BubbleInfo
+
+object FakeBubbleViewFactory {
+
+    /** Inflates a [BubbleView] and adds it to the [parent] view if it is present. */
+    fun createBubble(
+        context: Context,
+        key: String,
+        parent: ViewGroup?,
+        iconSize: Int = 50,
+        iconColor: Int,
+        badgeColor: Int = Color.RED,
+        dotColor: Int = Color.BLUE,
+        suppressNotification: Boolean = false,
+    ): BubbleView {
+        val inflater = LayoutInflater.from(context)
+        // BubbleView uses launcher's badge to icon ratio and expects the badge image to already
+        // have the right size
+        val badgeToIconRatio = 0.444f
+        val badgeRadius = iconSize * badgeToIconRatio / 2
+        val icon = createCircleBitmap(radius = iconSize / 2, color = iconColor)
+        val badge = createCircleBitmap(radius = badgeRadius.toInt(), color = badgeColor)
+
+        val flags =
+            if (suppressNotification) Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION else 0
+        val bubbleInfo =
+            BubbleInfo(key, flags, null, null, 0, context.packageName, null, null, false, true)
+        val bubbleView = inflater.inflate(R.layout.bubblebar_item_view, parent, false) as BubbleView
+        val dotPath =
+            PathParser.createPathFromPathData(
+                context.resources.getString(com.android.internal.R.string.config_icon_mask)
+            )
+        val bubble =
+            BubbleBarBubble(bubbleInfo, bubbleView, badge, icon, dotColor, dotPath, "test app")
+        bubbleView.setBubble(bubble)
+        return bubbleView
+    }
+
+    private fun createCircleBitmap(radius: Int, color: Int): Bitmap {
+        val bitmap = Bitmap.createBitmap(radius * 2, radius * 2, Bitmap.Config.ARGB_8888)
+        val canvas = Canvas(bitmap)
+        canvas.drawARGB(0, 0, 0, 0)
+        val paint = Paint()
+        paint.color = color
+        canvas.drawCircle(radius.toFloat(), radius.toFloat(), radius.toFloat(), paint)
+        return bitmap
+    }
+}
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt
new file mode 100644
index 0000000..e4b8069
--- /dev/null
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewScreenshotTest.kt
@@ -0,0 +1,141 @@
+/*
+ * 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.taskbar.bubbles
+
+import android.content.Context
+import android.graphics.Color
+import android.platform.test.rule.ScreenRecordRule
+import android.view.View
+import android.widget.FrameLayout
+import android.widget.FrameLayout.LayoutParams.MATCH_PARENT
+import android.widget.FrameLayout.LayoutParams.WRAP_CONTENT
+import androidx.activity.ComponentActivity
+import androidx.test.core.app.ApplicationProvider
+import com.android.launcher3.R
+import com.android.launcher3.taskbar.bubbles.testing.FakeBubbleViewFactory
+import com.google.android.apps.nexuslauncher.imagecomparison.goldenpathmanager.ViewScreenshotGoldenPathManager
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
+import platform.test.screenshot.DeviceEmulationSpec
+import platform.test.screenshot.Displays
+import platform.test.screenshot.ViewScreenshotTestRule
+import platform.test.screenshot.getEmulatedDevicePathConfig
+
+/** Screenshot tests for [BubbleBarView]. */
+@RunWith(ParameterizedAndroidJunit4::class)
+@ScreenRecordRule.ScreenRecord
+class BubbleBarViewScreenshotTest(emulationSpec: DeviceEmulationSpec) {
+
+    private val context = ApplicationProvider.getApplicationContext<Context>()
+    private lateinit var bubbleBarView: BubbleBarView
+
+    companion object {
+        @Parameters(name = "{0}")
+        @JvmStatic
+        fun getTestSpecs() =
+            DeviceEmulationSpec.forDisplays(
+                Displays.Phone,
+                isDarkTheme = false,
+                isLandscape = false,
+            )
+    }
+
+    @get:Rule
+    val screenshotRule =
+        ViewScreenshotTestRule(
+            emulationSpec,
+            ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec)),
+        )
+
+    @Test
+    fun bubbleBarView_collapsed_oneBubble() {
+        screenshotRule.screenshotTest("bubbleBarView_collapsed_oneBubble") { activity ->
+            activity.actionBar?.hide()
+            setupBubbleBarView()
+            bubbleBarView.addBubble(createBubble("key1", Color.GREEN))
+            val container = FrameLayout(context)
+            val lp = FrameLayout.LayoutParams(MATCH_PARENT, MATCH_PARENT)
+            container.layoutParams = lp
+            container.addView(bubbleBarView)
+            container
+        }
+    }
+
+    @Test
+    fun bubbleBarView_collapsed_twoBubbles() {
+        screenshotRule.screenshotTest("bubbleBarView_collapsed_twoBubbles") { activity ->
+            activity.actionBar?.hide()
+            setupBubbleBarView()
+            bubbleBarView.addBubble(createBubble("key1", Color.GREEN))
+            bubbleBarView.addBubble(createBubble("key2", Color.CYAN))
+            val container = FrameLayout(context)
+            val lp = FrameLayout.LayoutParams(MATCH_PARENT, MATCH_PARENT)
+            container.layoutParams = lp
+            container.addView(bubbleBarView)
+            container
+        }
+    }
+
+    @Test
+    fun bubbleBarView_expanded_threeBubbles() {
+        // if we're still expanding, wait with taking a screenshot
+        val shouldWait: (ComponentActivity, View) -> Boolean = { _, _ -> bubbleBarView.isExpanding }
+        screenshotRule.screenshotTest(
+            "bubbleBarView_expanded_threeBubbles",
+            checkView = shouldWait,
+        ) { activity ->
+            activity.actionBar?.hide()
+            setupBubbleBarView()
+            bubbleBarView.addBubble(createBubble("key1", Color.GREEN))
+            bubbleBarView.addBubble(createBubble("key2", Color.CYAN))
+            bubbleBarView.addBubble(createBubble("key3", Color.MAGENTA))
+            val container = FrameLayout(context)
+            val lp = FrameLayout.LayoutParams(MATCH_PARENT, MATCH_PARENT)
+            container.layoutParams = lp
+            container.addView(bubbleBarView)
+            bubbleBarView.isExpanded = true
+            container
+        }
+    }
+
+    private fun setupBubbleBarView() {
+        bubbleBarView = BubbleBarView(context)
+        val lp = FrameLayout.LayoutParams(WRAP_CONTENT, WRAP_CONTENT)
+        bubbleBarView.layoutParams = lp
+        val paddingTop =
+            context.resources.getDimensionPixelSize(R.dimen.bubblebar_pointer_visible_size)
+        bubbleBarView.setPadding(0, paddingTop, 0, 0)
+        bubbleBarView.visibility = View.VISIBLE
+        bubbleBarView.alpha = 1f
+    }
+
+    private fun createBubble(key: String, color: Int): BubbleView {
+        val bubbleView =
+            FakeBubbleViewFactory.createBubble(
+                context,
+                key,
+                parent = bubbleBarView,
+                iconColor = color,
+            )
+        bubbleView.showDotIfNeeded(1f)
+        bubbleBarView.setSelectedBubble(bubbleView)
+        return bubbleView
+    }
+}
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt
index eb459ff..47f393f 100644
--- a/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt
@@ -15,17 +15,10 @@
  */
 package com.android.launcher3.taskbar.bubbles
 
-import android.app.Notification
 import android.content.Context
-import android.graphics.Bitmap
-import android.graphics.Canvas
 import android.graphics.Color
-import android.graphics.Paint
-import android.util.PathParser
-import android.view.LayoutInflater
 import androidx.test.core.app.ApplicationProvider
-import com.android.launcher3.R
-import com.android.wm.shell.shared.bubbles.BubbleInfo
+import com.android.launcher3.taskbar.bubbles.testing.FakeBubbleViewFactory
 import com.google.android.apps.nexuslauncher.imagecomparison.goldenpathmanager.ViewScreenshotGoldenPathManager
 import org.junit.Rule
 import org.junit.Test
@@ -50,7 +43,7 @@
             DeviceEmulationSpec.forDisplays(
                 Displays.Phone,
                 isDarkTheme = false,
-                isLandscape = false
+                isLandscape = false,
             )
     }
 
@@ -58,7 +51,7 @@
     val screenshotRule =
         ViewScreenshotTestRule(
             emulationSpec,
-            ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec))
+            ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec)),
         )
 
     @Test
@@ -86,39 +79,16 @@
     }
 
     private fun setupBubbleView(suppressNotification: Boolean = false): BubbleView {
-        val inflater = LayoutInflater.from(context)
-
-        val iconSize = 100
-        // BubbleView uses launcher's badge to icon ratio and expects the badge image to already
-        // have the right size
-        val badgeToIconRatio = 0.444f
-        val badgeRadius = iconSize * badgeToIconRatio / 2
-        val icon = createCircleBitmap(radius = iconSize / 2, color = Color.LTGRAY)
-        val badge = createCircleBitmap(radius = badgeRadius.toInt(), color = Color.RED)
-
-        val flags =
-            if (suppressNotification) Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION else 0
-        val bubbleInfo =
-            BubbleInfo("key", flags, null, null, 0, context.packageName, null, null, false, true)
-        val bubbleView = inflater.inflate(R.layout.bubblebar_item_view, null) as BubbleView
-        val dotPath =
-            PathParser.createPathFromPathData(
-                context.resources.getString(com.android.internal.R.string.config_icon_mask)
+        val bubbleView =
+            FakeBubbleViewFactory.createBubble(
+                context,
+                key = "key",
+                parent = null,
+                iconSize = 100,
+                iconColor = Color.LTGRAY,
+                suppressNotification = suppressNotification,
             )
-        val bubble =
-            BubbleBarBubble(bubbleInfo, bubbleView, badge, icon, Color.BLUE, dotPath, "test app")
-        bubbleView.setBubble(bubble)
         bubbleView.showDotIfNeeded(1f)
         return bubbleView
     }
-
-    private fun createCircleBitmap(radius: Int, color: Int): Bitmap {
-        val bitmap = Bitmap.createBitmap(radius * 2, radius * 2, Bitmap.Config.ARGB_8888)
-        val canvas = Canvas(bitmap)
-        canvas.drawARGB(0, 0, 0, 0)
-        val paint = Paint()
-        paint.color = color
-        canvas.drawCircle(radius.toFloat(), radius.toFloat(), radius.toFloat(), paint)
-        return bitmap
-    }
 }