feat: Add option to clear home screen in settings (#6125)

Signed-off-by: abhixv <abhi.sharma1@hotmail.com>

quickstep/tests/multivalentTests/src/com/android/launcher3/model/AppEventProducerTest.java

@@ -35,13 +35,10 @@ import android.os.UserHandle;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
-import com.android.launcher3.dagger.LauncherAppComponent;
-import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.util.AllModulesForTest;
-import com.android.launcher3.util.SandboxContext;
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 import com.android.launcher3.util.UserIconInfo;
 import com.android.systemui.shared.system.SysUiStatsLog;
 
@@ -54,9 +51,6 @@ import org.mockito.MockitoAnnotations;
 
 import java.util.Arrays;
 
-import dagger.BindsInstance;
-import dagger.Component;
-
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class AppEventProducerTest {
@@ -78,9 +72,7 @@ public class AppEventProducerTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mContext = new SandboxContext(getApplicationContext());
-        mContext.initDaggerComponent(
-                DaggerAppEventProducerTest_TestComponent.builder().bindUserCache(mUserCache)
-        );
+        mContext.putObject(UserCache.INSTANCE, mUserCache);
         mAppEventProducer = new AppEventProducer(mContext, null);
     }
 
@@ -137,15 +129,4 @@ public class AppEventProducerTest {
                 .build());
         return itemBuilder.build();
     }
-
-    @LauncherAppSingleton
-    @Component(modules = { AllModulesForTest.class })
-    interface TestComponent extends LauncherAppComponent {
-        @Component.Builder
-        interface Builder extends LauncherAppComponent.Builder {
-            @BindsInstance
-            AppEventProducerTest.TestComponent.Builder bindUserCache(UserCache userCache);
-            @Override LauncherAppComponent build();
-        }
-    }
 }

quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java

@@ -4,8 +4,6 @@ import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCH
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_BACK_BUTTON_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_HOME_BUTTON_LONGPRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_HOME_BUTTON_TAP;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_LONGPRESS;
-import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_OVERVIEW_BUTTON_LONGPRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_OVERVIEW_BUTTON_TAP;
 import static com.android.launcher3.taskbar.TaskbarNavButtonController.BUTTON_A11Y;
@@ -15,14 +13,11 @@ import static com.android.launcher3.taskbar.TaskbarNavButtonController.BUTTON_IM
 import static com.android.launcher3.taskbar.TaskbarNavButtonController.BUTTON_RECENTS;
 import static com.android.launcher3.taskbar.TaskbarNavButtonController.SCREEN_PIN_LONG_PRESS_THRESHOLD;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
-import static com.android.window.flags.Flags.FLAG_PREDICTIVE_BACK_THREE_BUTTON_NAV;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -30,12 +25,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.os.Handler;
-import android.platform.test.annotations.RequiresFlagsEnabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
-import android.view.KeyEvent;
 import android.view.View;
-import android.view.inputmethod.Flags;
 
 import androidx.test.platform.app.InstrumentationRegistry;
 import androidx.test.runner.AndroidJUnit4;
@@ -44,14 +34,11 @@ import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TouchInteractionService;
-import com.android.quickstep.util.ContextualSearchInvoker;
-import com.android.systemui.contextualeducation.GestureType;
+import com.android.quickstep.util.AssistUtils;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -62,13 +49,12 @@ public class TaskbarNavButtonControllerTest {
 
     @Mock
     SystemUiProxy mockSystemUiProxy;
-
     @Mock
     TouchInteractionService mockService;
     @Mock
     Handler mockHandler;
     @Mock
-    ContextualSearchInvoker mockContextualSearchInvoker;
+    AssistUtils mockAssistUtils;
     @Mock
     StatsLogManager mockStatsLogManager;
     @Mock
@@ -80,9 +66,6 @@ public class TaskbarNavButtonControllerTest {
     @Mock
     View mockView;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
-
     private int mHomePressCount;
     private int mOverviewToggleCount;
     private final TaskbarNavButtonCallbacks mCallbacks = new TaskbarNavButtonCallbacks() {
@@ -115,45 +98,19 @@ public class TaskbarNavButtonControllerTest {
                 mCallbacks,
                 mockSystemUiProxy,
                 mockHandler,
-                mockContextualSearchInvoker);
+                mockAssistUtils);
     }
 
     @Test
     public void testPressBack() {
         mNavButtonController.onButtonClick(BUTTON_BACK, mockView);
-        verify(mockSystemUiProxy, times(1)).onBackEvent(null);
-    }
-
-    @Test
-    public void testPressBack_updateContextualEduData() {
-        mNavButtonController.onButtonClick(BUTTON_BACK, mockView);
-        verify(mockSystemUiProxy, times(1))
-                .updateContextualEduStats(/* isTrackpad= */ eq(false), eq(GestureType.BACK));
+        verify(mockSystemUiProxy, times(1)).onBackPressed();
     }
 
     @Test
     public void testPressImeSwitcher() {
-        mNavButtonController.init(mockTaskbarControllers);
         mNavButtonController.onButtonClick(BUTTON_IME_SWITCH, mockView);
-        verify(mockStatsLogger, times(1)).log(LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_TAP);
-        verify(mockStatsLogger, never()).log(LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_LONGPRESS);
         verify(mockSystemUiProxy, times(1)).onImeSwitcherPressed();
-        verify(mockSystemUiProxy, never()).onImeSwitcherLongPress();
-    }
-
-    @Test
-    public void testLongPressImeSwitcher() {
-        mNavButtonController.init(mockTaskbarControllers);
-        mNavButtonController.onButtonLongClick(BUTTON_IME_SWITCH, mockView);
-        verify(mockStatsLogger, never()).log(LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_TAP);
-        verify(mockSystemUiProxy, never()).onImeSwitcherPressed();
-        if (Flags.imeSwitcherRevamp()) {
-            verify(mockStatsLogger, times(1)).log(LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_LONGPRESS);
-            verify(mockSystemUiProxy, times(1)).onImeSwitcherLongPress();
-        } else {
-            verify(mockStatsLogger, never()).log(LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_LONGPRESS);
-            verify(mockSystemUiProxy, never()).onImeSwitcherLongPress();
-        }
     }
 
     @Test
@@ -172,40 +129,40 @@ public class TaskbarNavButtonControllerTest {
     @Test
     public void testLongPressHome_enabled_withoutOverride() {
         mNavButtonController.setAssistantLongPressEnabled(true /*assistantLongPressEnabled*/);
-        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(false);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(false);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockContextualSearchInvoker, times(1)).tryStartAssistOverride(anyInt());
+        verify(mockAssistUtils, times(1)).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, times(1)).startAssistant(any());
     }
 
     @Test
     public void testLongPressHome_enabled_withOverride() {
         mNavButtonController.setAssistantLongPressEnabled(true /*assistantLongPressEnabled*/);
-        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(true);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(true);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockContextualSearchInvoker, times(1)).tryStartAssistOverride(anyInt());
+        verify(mockAssistUtils, times(1)).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
     @Test
     public void testLongPressHome_disabled_withoutOverride() {
         mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
-        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(false);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(false);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockContextualSearchInvoker, never()).tryStartAssistOverride(anyInt());
+        verify(mockAssistUtils, never()).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
     @Test
     public void testLongPressHome_disabled_withOverride() {
         mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
-        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(true);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(true);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockContextualSearchInvoker, never()).tryStartAssistOverride(anyInt());
+        verify(mockAssistUtils, never()).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
@@ -215,26 +172,12 @@ public class TaskbarNavButtonControllerTest {
         assertThat(mHomePressCount).isEqualTo(1);
     }
 
-    @Test
-    public void testPressHome_updateContextualEduData() {
-        mNavButtonController.onButtonClick(BUTTON_HOME, mockView);
-        verify(mockSystemUiProxy, times(1))
-                .updateContextualEduStats(/* isTrackpad= */ eq(false), eq(GestureType.HOME));
-    }
-
     @Test
     public void testPressRecents() {
         mNavButtonController.onButtonClick(BUTTON_RECENTS, mockView);
         assertThat(mOverviewToggleCount).isEqualTo(1);
     }
 
-    @Test
-    public void testPressRecents_updateContextualEduData() {
-        mNavButtonController.onButtonClick(BUTTON_RECENTS, mockView);
-        verify(mockSystemUiProxy, times(1))
-                .updateContextualEduStats(/* isTrackpad= */ eq(false), eq(GestureType.OVERVIEW));
-    }
-
     @Test
     public void testPressRecentsWithScreenPinned_noNavigationToOverview() {
         mNavButtonController.updateSysuiFlags(SYSUI_STATE_SCREEN_PINNING);
@@ -339,46 +282,4 @@ public class TaskbarNavButtonControllerTest {
         verify(mockStatsLogger, times(1)).log(LAUNCHER_TASKBAR_BACK_BUTTON_LONGPRESS);
         verify(mockStatsLogger, times(0)).log(LAUNCHER_TASKBAR_BACK_BUTTON_TAP);
     }
-
-    @Test
-    @RequiresFlagsEnabled(FLAG_PREDICTIVE_BACK_THREE_BUTTON_NAV)
-    public void testPredictiveBackInvoked() {
-        ArgumentCaptor<KeyEvent> keyEventCaptor = ArgumentCaptor.forClass(KeyEvent.class);
-        mNavButtonController.sendBackKeyEvent(KeyEvent.ACTION_DOWN, false);
-        mNavButtonController.sendBackKeyEvent(KeyEvent.ACTION_UP, false);
-        verify(mockSystemUiProxy, times(2)).onBackEvent(keyEventCaptor.capture());
-        verifyKeyEvent(keyEventCaptor.getAllValues().getFirst(), KeyEvent.ACTION_DOWN, false);
-        verifyKeyEvent(keyEventCaptor.getAllValues().getLast(), KeyEvent.ACTION_UP, false);
-    }
-
-    @Test
-    @RequiresFlagsEnabled(FLAG_PREDICTIVE_BACK_THREE_BUTTON_NAV)
-    public void testPredictiveBackCancelled() {
-        ArgumentCaptor<KeyEvent> keyEventCaptor = ArgumentCaptor.forClass(KeyEvent.class);
-        mNavButtonController.sendBackKeyEvent(KeyEvent.ACTION_DOWN, false);
-        mNavButtonController.sendBackKeyEvent(KeyEvent.ACTION_UP, true);
-        verify(mockSystemUiProxy, times(2)).onBackEvent(keyEventCaptor.capture());
-        verifyKeyEvent(keyEventCaptor.getAllValues().getFirst(), KeyEvent.ACTION_DOWN, false);
-        verifyKeyEvent(keyEventCaptor.getAllValues().getLast(), KeyEvent.ACTION_UP, true);
-    }
-
-    @Test
-    @RequiresFlagsEnabled(FLAG_PREDICTIVE_BACK_THREE_BUTTON_NAV)
-    public void testButtonsDisabledWhileBackPressed() {
-        mNavButtonController.sendBackKeyEvent(KeyEvent.ACTION_DOWN, false);
-        mNavButtonController.onButtonClick(BUTTON_HOME, mockView);
-        mNavButtonController.onButtonClick(BUTTON_RECENTS, mockView);
-        mNavButtonController.onButtonLongClick(BUTTON_A11Y, mockView);
-        mNavButtonController.onButtonClick(BUTTON_IME_SWITCH, mockView);
-        mNavButtonController.sendBackKeyEvent(KeyEvent.ACTION_UP, false);
-        assertThat(mHomePressCount).isEqualTo(0);
-        verify(mockSystemUiProxy, never()).notifyAccessibilityButtonLongClicked();
-        assertThat(mOverviewToggleCount).isEqualTo(0);
-        verify(mockSystemUiProxy, never()).onImeSwitcherPressed();
-    }
-
-    private void verifyKeyEvent(KeyEvent keyEvent, int action, boolean isCancelled) {
-        assertEquals(isCancelled, keyEvent.isCanceled());
-        assertEquals(action, KeyEvent.ACTION_DOWN, keyEvent.getAction());
-    }
 }

quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt

@@ -0,0 +1,164 @@
+/*
+ * 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
+
+import android.app.Instrumentation
+import android.app.PendingIntent
+import android.content.IIntentSender
+import android.content.Intent
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.rule.ServiceTestRule
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks
+import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
+import com.android.launcher3.util.LauncherMultivalentJUnit.Companion.isRunningInRobolectric
+import com.android.quickstep.AllAppsActionManager
+import com.android.quickstep.TouchInteractionService
+import com.android.quickstep.TouchInteractionService.TISBinder
+import org.junit.Assume.assumeTrue
+import org.junit.rules.MethodRule
+import org.junit.runners.model.FrameworkMethod
+import org.junit.runners.model.Statement
+
+/**
+ * Manages the Taskbar lifecycle for unit tests.
+ *
+ * See [InjectController] for grabbing controller(s) under test with minimal boilerplate.
+ *
+ * The rule interacts with [TaskbarManager] on the main thread. A good rule of thumb for tests is
+ * that code that is executed on the main thread in production should also happen on that thread
+ * when tested.
+ *
+ * `@UiThreadTest` is a simple way to run an entire test body on the main thread. But if a test
+ * executes code that appends message(s) to the main thread's `MessageQueue`, the annotation will
+ * prevent those messages from being processed until after the test body finishes.
+ *
+ * To test pending messages, instead use something like [Instrumentation.runOnMainSync] to perform
+ * only sections of the test body on the main thread synchronously:
+ * ```
+ * @Test
+ * fun example() {
+ *     instrumentation.runOnMainSync { doWorkThatPostsMessage() }
+ *     // Second lambda will not execute until message is processed.
+ *     instrumentation.runOnMainSync { verifyMessageResults() }
+ * }
+ * ```
+ */
+class TaskbarUnitTestRule : MethodRule {
+    private val instrumentation = InstrumentationRegistry.getInstrumentation()
+    private val serviceTestRule = ServiceTestRule()
+
+    private lateinit var taskbarManager: TaskbarManager
+    private lateinit var target: Any
+
+    val activityContext: TaskbarActivityContext
+        get() {
+            return taskbarManager.currentActivityContext
+                ?: throw RuntimeException("Failed to obtain TaskbarActivityContext.")
+        }
+
+    override fun apply(base: Statement, method: FrameworkMethod, target: Any): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                this@TaskbarUnitTestRule.target = target
+
+                val context = instrumentation.targetContext
+                instrumentation.runOnMainSync {
+                    assumeTrue(
+                        LauncherAppState.getIDP(context).getDeviceProfile(context).isTaskbarPresent
+                    )
+                }
+
+                // Check for existing Taskbar instance from Launcher process.
+                val launcherTaskbarManager: TaskbarManager? =
+                    if (!isRunningInRobolectric) {
+                        try {
+                            val tisBinder =
+                                serviceTestRule.bindService(
+                                    Intent(context, TouchInteractionService::class.java)
+                                ) as? TISBinder
+                            tisBinder?.taskbarManager
+                        } catch (_: Exception) {
+                            null
+                        }
+                    } else {
+                        null
+                    }
+
+                instrumentation.runOnMainSync {
+                    taskbarManager =
+                        TaskbarManager(
+                            context,
+                            AllAppsActionManager(context, UI_HELPER_EXECUTOR) {
+                                PendingIntent(IIntentSender.Default())
+                            },
+                            object : TaskbarNavButtonCallbacks {},
+                        )
+                }
+
+                try {
+                    // Replace Launcher Taskbar window with test instance.
+                    instrumentation.runOnMainSync {
+                        launcherTaskbarManager?.removeTaskbarRootViewFromWindow()
+                        taskbarManager.onUserUnlocked() // Required to complete initialization.
+                    }
+
+                    injectControllers()
+                    base.evaluate()
+                } finally {
+                    // Revert Taskbar window.
+                    instrumentation.runOnMainSync {
+                        taskbarManager.destroy()
+                        launcherTaskbarManager?.addTaskbarRootViewToWindow()
+                    }
+                }
+            }
+        }
+    }
+
+    /** Simulates Taskbar recreation lifecycle. */
+    fun recreateTaskbar() {
+        taskbarManager.recreateTaskbar()
+        injectControllers()
+    }
+
+    private fun injectControllers() {
+        val controllers = activityContext.controllers
+        val controllerFieldsByType = controllers.javaClass.fields.associateBy { it.type }
+        target.javaClass.fields
+            .filter { it.isAnnotationPresent(InjectController::class.java) }
+            .forEach {
+                it.set(
+                    target,
+                    controllerFieldsByType[it.type]?.get(controllers)
+                        ?: throw NoSuchElementException("Failed to find controller for ${it.type}"),
+                )
+            }
+    }
+
+    /**
+     * Annotates test controller fields to inject the corresponding controllers from the current
+     * [TaskbarControllers] instance.
+     *
+     * Controllers are injected during test setup and upon calling [recreateTaskbar].
+     *
+     * Multiple controllers can be injected if needed.
+     */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.FIELD)
+    annotation class InjectController
+}

quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt

@@ -16,25 +16,22 @@
 
 package com.android.launcher3.taskbar.allapps
 
-import android.animation.AnimatorTestRule
 import android.content.ComponentName
 import android.content.Intent
 import android.os.Process
+import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.BubbleTextView
 import com.android.launcher3.appprediction.PredictionRowView
 import com.android.launcher3.model.data.AppInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.notification.NotificationKeyData
-import com.android.launcher3.taskbar.TaskbarControllerTestUtil.runOnMainSync
+import com.android.launcher3.taskbar.TaskbarUnitTestRule
+import com.android.launcher3.taskbar.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayController
-import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
-import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
-import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.util.PackageUserKey
-import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
 import org.junit.Test
@@ -44,98 +41,89 @@ import org.junit.runner.RunWith
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarAllAppsControllerTest {
 
-    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
-    @get:Rule(order = 1) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
-    @get:Rule(order = 2) val animatorTestRule = AnimatorTestRule(this)
+    @get:Rule val taskbarUnitTestRule = TaskbarUnitTestRule()
 
     @InjectController lateinit var allAppsController: TaskbarAllAppsController
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     @Test
+    @UiThreadTest
     fun testToggle_once_showsAllApps() {
-        runOnMainSync { allAppsController.toggle() }
+        allAppsController.toggle()
         assertThat(allAppsController.isOpen).isTrue()
     }
 
     @Test
+    @UiThreadTest
     fun testToggle_twice_closesAllApps() {
-        runOnMainSync {
-            allAppsController.toggle()
-            allAppsController.toggle()
-        }
+        allAppsController.toggle()
+        allAppsController.toggle()
         assertThat(allAppsController.isOpen).isFalse()
     }
 
     @Test
+    @UiThreadTest
     fun testToggle_taskbarRecreated_allAppsReopened() {
-        runOnMainSync { allAppsController.toggle() }
+        allAppsController.toggle()
         taskbarUnitTestRule.recreateTaskbar()
         assertThat(allAppsController.isOpen).isTrue()
     }
 
     @Test
+    @UiThreadTest
     fun testSetApps_beforeOpened_cachesInfo() {
-        val overlayContext =
-            TestUtil.getOnUiThread {
-                allAppsController.setApps(TEST_APPS, 0, emptyMap())
-                allAppsController.toggle()
-                overlayController.requestWindow()
-            }
+        allAppsController.setApps(TEST_APPS, 0, emptyMap())
+        allAppsController.toggle()
 
+        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
+    @UiThreadTest
     fun testSetApps_afterOpened_updatesStore() {
-        val overlayContext =
-            TestUtil.getOnUiThread {
-                allAppsController.toggle()
-                allAppsController.setApps(TEST_APPS, 0, emptyMap())
-                overlayController.requestWindow()
-            }
+        allAppsController.toggle()
+        allAppsController.setApps(TEST_APPS, 0, emptyMap())
 
+        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
+    @UiThreadTest
     fun testSetPredictedApps_beforeOpened_cachesInfo() {
+        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+        allAppsController.toggle()
+
         val predictedApps =
-            TestUtil.getOnUiThread {
-                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-                allAppsController.toggle()
-
-                overlayController
-                    .requestWindow()
-                    .appsView
-                    .floatingHeaderView
-                    .findFixedRowByType(PredictionRowView::class.java)
-                    .predictedApps
-            }
-
+            overlayController
+                .requestWindow()
+                .appsView
+                .floatingHeaderView
+                .findFixedRowByType(PredictionRowView::class.java)
+                .predictedApps
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
     @Test
+    @UiThreadTest
     fun testSetPredictedApps_afterOpened_cachesInfo() {
+        allAppsController.toggle()
+        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+
         val predictedApps =
-            TestUtil.getOnUiThread {
-                allAppsController.toggle()
-                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-
-                overlayController
-                    .requestWindow()
-                    .appsView
-                    .floatingHeaderView
-                    .findFixedRowByType(PredictionRowView::class.java)
-                    .predictedApps
-            }
-
+            overlayController
+                .requestWindow()
+                .appsView
+                .floatingHeaderView
+                .findFixedRowByType(PredictionRowView::class.java)
+                .predictedApps
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
     @Test
     fun testUpdateNotificationDots_appInfo_hasDot() {
-        runOnMainSync {
+        getInstrumentation().runOnMainSync {
             allAppsController.setApps(TEST_APPS, 0, emptyMap())
             allAppsController.toggle()
             taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
@@ -145,58 +133,41 @@ class TaskbarAllAppsControllerTest {
         }
 
         // Ensure the recycler view fully inflates before trying to grab an icon.
-        val btv =
-            TestUtil.getOnUiThread {
+        getInstrumentation().runOnMainSync {
+            val btv =
                 overlayController
                     .requestWindow()
                     .appsView
                     .activeRecyclerView
                     .findViewHolderForAdapterPosition(0)
                     ?.itemView as? BubbleTextView
-            }
-        assertThat(btv?.hasDot()).isTrue()
+            assertThat(btv?.hasDot()).isTrue()
+        }
     }
 
     @Test
+    @UiThreadTest
     fun testUpdateNotificationDots_predictedApp_hasDot() {
-        runOnMainSync {
-            allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-            allAppsController.toggle()
-            taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
-                PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
-                NotificationKeyData("key"),
-            )
-        }
+        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+        allAppsController.toggle()
 
-        val btv =
-            TestUtil.getOnUiThread {
-                overlayController
-                    .requestWindow()
-                    .appsView
-                    .floatingHeaderView
-                    .findFixedRowByType(PredictionRowView::class.java)
-                    .getChildAt(0) as BubbleTextView
-            }
+        taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
+            PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
+            NotificationKeyData("key"),
+        )
+
+        val predictionRowView =
+            overlayController
+                .requestWindow()
+                .appsView
+                .floatingHeaderView
+                .findFixedRowByType(PredictionRowView::class.java)
+        val btv = predictionRowView.getChildAt(0) as BubbleTextView
         assertThat(btv.hasDot()).isTrue()
     }
 
-    @Test
-    fun testToggleSearch_searchEditTextFocused() {
-        runOnMainSync { allAppsController.toggleSearch() }
-        runOnMainSync {
-            // All Apps is now attached to window. Open animation is posted but not started.
-        }
-
-        runOnMainSync {
-            // Animation has started. Advance to end of animation.
-            animatorTestRule.advanceTimeBy(overlayController.openDuration.toLong())
-        }
-        val editText = overlayController.requestWindow().appsView.searchUiManager.editText
-        assertThat(editText?.hasFocus()).isTrue()
-    }
-
-    companion object {
-        val TEST_APPS =
+    private companion object {
+        private val TEST_APPS =
             Array(16) {
                 AppInfo(
                     ComponentName(
@@ -209,6 +180,6 @@ class TaskbarAllAppsControllerTest {
                 )
             }
 
-        val TEST_PREDICTED_APPS = TEST_APPS.take(4).map { WorkspaceItemInfo(it) }
+        private val TEST_PREDICTED_APPS = TEST_APPS.take(4).map { WorkspaceItemInfo(it) }
     }
 }

quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt

@@ -0,0 +1,454 @@
+/*
+ * 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.animation
+
+import android.content.Context
+import android.graphics.Color
+import android.graphics.Path
+import android.graphics.drawable.ColorDrawable
+import android.view.LayoutInflater
+import android.view.View
+import android.view.View.VISIBLE
+import android.widget.FrameLayout
+import androidx.core.graphics.drawable.toBitmap
+import androidx.dynamicanimation.animation.DynamicAnimation
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.R
+import com.android.launcher3.taskbar.bubbles.BubbleBarBubble
+import com.android.launcher3.taskbar.bubbles.BubbleBarOverflow
+import com.android.launcher3.taskbar.bubbles.BubbleBarView
+import com.android.launcher3.taskbar.bubbles.BubbleStashController
+import com.android.launcher3.taskbar.bubbles.BubbleView
+import com.android.wm.shell.common.bubbles.BubbleInfo
+import com.android.wm.shell.shared.animation.PhysicsAnimator
+import com.android.wm.shell.shared.animation.PhysicsAnimatorTestUtils
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.atLeastOnce
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class BubbleBarViewAnimatorTest {
+
+    private val context = ApplicationProvider.getApplicationContext<Context>()
+    private lateinit var animatorScheduler: TestBubbleBarViewAnimatorScheduler
+    private lateinit var overflowView: BubbleView
+    private lateinit var bubbleView: BubbleView
+    private lateinit var bubble: BubbleBarBubble
+    private lateinit var bubbleBarView: BubbleBarView
+    private lateinit var bubbleStashController: BubbleStashController
+
+    @Before
+    fun setUp() {
+        animatorScheduler = TestBubbleBarViewAnimatorScheduler()
+        PhysicsAnimatorTestUtils.prepareForTest()
+    }
+
+    @Test
+    fun animateBubbleInForStashed() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateBubbleInForStashed(bubble)
+        }
+
+        // let the animation start and wait for it to complete
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        assertThat(handle.alpha).isEqualTo(0)
+        assertThat(handle.translationY)
+            .isEqualTo(DIFF_BETWEEN_HANDLE_AND_BAR_CENTERS + BAR_TRANSLATION_Y_FOR_TASKBAR)
+        assertThat(bubbleBarView.visibility).isEqualTo(VISIBLE)
+        assertThat(bubbleBarView.scaleX).isEqualTo(1)
+        assertThat(bubbleBarView.scaleY).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_TASKBAR)
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+
+        // execute the hide bubble animation
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(animatorScheduler.delayedBlock!!)
+
+        // let the animation start and wait for it to complete
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        assertThat(handle.alpha).isEqualTo(1)
+        assertThat(handle.translationY).isEqualTo(0)
+        assertThat(bubbleBarView.alpha).isEqualTo(0)
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        verify(bubbleStashController).stashBubbleBarImmediate()
+    }
+
+    @Test
+    fun animateBubbleInForStashed_tapAnimatingBubble() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateBubbleInForStashed(bubble)
+        }
+
+        // let the animation start and wait for it to complete
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        assertThat(handle.alpha).isEqualTo(0)
+        assertThat(handle.translationY)
+            .isEqualTo(DIFF_BETWEEN_HANDLE_AND_BAR_CENTERS + BAR_TRANSLATION_Y_FOR_TASKBAR)
+        assertThat(bubbleBarView.visibility).isEqualTo(VISIBLE)
+        assertThat(bubbleBarView.scaleX).isEqualTo(1)
+        assertThat(bubbleBarView.scaleY).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_TASKBAR)
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+
+        verify(bubbleStashController, atLeastOnce()).updateTaskbarTouchRegion()
+
+        // verify the hide bubble animation is pending
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+
+        animator.onBubbleBarTouchedWhileAnimating()
+
+        assertThat(animatorScheduler.delayedBlock).isNull()
+        assertThat(bubbleBarView.alpha).isEqualTo(1)
+        assertThat(bubbleBarView.visibility).isEqualTo(VISIBLE)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_TASKBAR)
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+    }
+
+    @Test
+    fun animateBubbleInForStashed_touchTaskbarArea_whileShowing() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateBubbleInForStashed(bubble)
+        }
+
+        // wait for the animation to start
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilFirstAnimationFrameWhereTrue(handleAnimator) { true }
+
+        handleAnimator.assertIsRunning()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+        // verify the hide bubble animation is pending
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.onStashStateChangingWhileAnimating()
+        }
+
+        // verify that the hide animation was canceled
+        assertThat(animatorScheduler.delayedBlock).isNull()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        verify(bubbleStashController).onNewBubbleAnimationInterrupted(any(), any())
+
+        // PhysicsAnimatorTestUtils posts the cancellation to the main thread so we need to wait
+        // again
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+        handleAnimator.assertIsNotRunning()
+    }
+
+    @Test
+    fun animateBubbleInForStashed_touchTaskbarArea_whileHiding() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateBubbleInForStashed(bubble)
+        }
+
+        // let the animation start and wait for it to complete
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        // execute the hide bubble animation
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(animatorScheduler.delayedBlock!!)
+
+        // wait for the hide animation to start
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        handleAnimator.assertIsRunning()
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.onStashStateChangingWhileAnimating()
+        }
+
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        verify(bubbleStashController).onNewBubbleAnimationInterrupted(any(), any())
+
+        // PhysicsAnimatorTestUtils posts the cancellation to the main thread so we need to wait
+        // again
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+        handleAnimator.assertIsNotRunning()
+    }
+
+    @Test
+    fun animateBubbleInForStashed_showAnimationCanceled() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateBubbleInForStashed(bubble)
+        }
+
+        // wait for the animation to start
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilFirstAnimationFrameWhereTrue(handleAnimator) { true }
+
+        handleAnimator.assertIsRunning()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+
+        handleAnimator.cancel()
+        handleAnimator.assertIsNotRunning()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        assertThat(animatorScheduler.delayedBlock).isNull()
+    }
+
+    @Test
+    fun animateToInitialState_inApp() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+        whenever(bubbleStashController.bubbleBarTranslationY)
+            .thenReturn(BAR_TRANSLATION_Y_FOR_TASKBAR)
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val barAnimator = PhysicsAnimator.getInstance(bubbleBarView)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateToInitialState(bubble, isInApp = true, isExpanding = false)
+        }
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        barAnimator.assertIsNotRunning()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+        assertThat(bubbleBarView.alpha).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_TASKBAR)
+
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(animatorScheduler.delayedBlock!!)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        assertThat(bubbleBarView.alpha).isEqualTo(0)
+        assertThat(handle.translationY).isEqualTo(0)
+        assertThat(handle.alpha).isEqualTo(1)
+
+        verify(bubbleStashController).stashBubbleBarImmediate()
+    }
+
+    @Test
+    fun animateToInitialState_inApp_autoExpanding() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+        whenever(bubbleStashController.bubbleBarTranslationY)
+            .thenReturn(BAR_TRANSLATION_Y_FOR_TASKBAR)
+
+        val handle = View(context)
+        val handleAnimator = PhysicsAnimator.getInstance(handle)
+        whenever(bubbleStashController.stashedHandlePhysicsAnimator).thenReturn(handleAnimator)
+
+        val barAnimator = PhysicsAnimator.getInstance(bubbleBarView)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateToInitialState(bubble, isInApp = true, isExpanding = true)
+        }
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        barAnimator.assertIsNotRunning()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+        assertThat(bubbleBarView.alpha).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_TASKBAR)
+
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(animatorScheduler.delayedBlock!!)
+
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        assertThat(bubbleBarView.alpha).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_TASKBAR)
+
+        verify(bubbleStashController).showBubbleBarImmediate()
+    }
+
+    @Test
+    fun animateToInitialState_inHome() {
+        setUpBubbleBar()
+        setUpBubbleStashController()
+        whenever(bubbleStashController.bubbleBarTranslationY)
+            .thenReturn(BAR_TRANSLATION_Y_FOR_HOTSEAT)
+
+        val barAnimator = PhysicsAnimator.getInstance(bubbleBarView)
+
+        val animator =
+            BubbleBarViewAnimator(bubbleBarView, bubbleStashController, animatorScheduler)
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animator.animateToInitialState(bubble, isInApp = false, isExpanding = false)
+        }
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {}
+        PhysicsAnimatorTestUtils.blockUntilAnimationsEnd(DynamicAnimation.TRANSLATION_Y)
+
+        barAnimator.assertIsNotRunning()
+        assertThat(bubbleBarView.isAnimatingNewBubble).isTrue()
+        assertThat(bubbleBarView.alpha).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_HOTSEAT)
+
+        assertThat(animatorScheduler.delayedBlock).isNotNull()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(animatorScheduler.delayedBlock!!)
+
+        assertThat(bubbleBarView.isAnimatingNewBubble).isFalse()
+        assertThat(bubbleBarView.alpha).isEqualTo(1)
+        assertThat(bubbleBarView.translationY).isEqualTo(BAR_TRANSLATION_Y_FOR_HOTSEAT)
+
+        verify(bubbleStashController).showBubbleBarImmediate()
+    }
+
+    private fun setUpBubbleBar() {
+        bubbleBarView = BubbleBarView(context)
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            bubbleBarView.layoutParams = FrameLayout.LayoutParams(0, 0)
+            val inflater = LayoutInflater.from(context)
+
+            val bitmap = ColorDrawable(Color.WHITE).toBitmap(width = 20, height = 20)
+            overflowView =
+                inflater.inflate(R.layout.bubblebar_item_view, bubbleBarView, false) as BubbleView
+            overflowView.setOverflow(BubbleBarOverflow(overflowView), bitmap)
+            bubbleBarView.addView(overflowView)
+
+            val bubbleInfo =
+                BubbleInfo("key", 0, null, null, 0, context.packageName, null, null, false)
+            bubbleView =
+                inflater.inflate(R.layout.bubblebar_item_view, bubbleBarView, false) as BubbleView
+            bubble =
+                BubbleBarBubble(bubbleInfo, bubbleView, bitmap, bitmap, Color.WHITE, Path(), "")
+            bubbleView.setBubble(bubble)
+            bubbleBarView.addView(bubbleView)
+        }
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+    }
+
+    private fun setUpBubbleStashController() {
+        bubbleStashController = mock<BubbleStashController>()
+        whenever(bubbleStashController.isStashed).thenReturn(true)
+        whenever(bubbleStashController.diffBetweenHandleAndBarCenters)
+            .thenReturn(DIFF_BETWEEN_HANDLE_AND_BAR_CENTERS)
+        whenever(bubbleStashController.stashedHandleTranslationForNewBubbleAnimation)
+            .thenReturn(HANDLE_TRANSLATION)
+        whenever(bubbleStashController.bubbleBarTranslationYForTaskbar)
+            .thenReturn(BAR_TRANSLATION_Y_FOR_TASKBAR)
+    }
+
+    private fun <T> PhysicsAnimator<T>.assertIsRunning() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            assertThat(isRunning()).isTrue()
+        }
+    }
+
+    private fun <T> PhysicsAnimator<T>.assertIsNotRunning() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            assertThat(isRunning()).isFalse()
+        }
+    }
+
+    private class TestBubbleBarViewAnimatorScheduler : BubbleBarViewAnimator.Scheduler {
+
+        var delayedBlock: Runnable? = null
+            private set
+
+        override fun post(block: Runnable) {
+            block.run()
+        }
+
+        override fun postDelayed(delayMillis: Long, block: Runnable) {
+            check(delayedBlock == null) { "there is already a pending block waiting to run" }
+            delayedBlock = block
+        }
+
+        override fun cancel(block: Runnable) {
+            check(delayedBlock == block) { "the pending block does not match the canceled block" }
+            delayedBlock = null
+        }
+    }
+}
+
+private const val DIFF_BETWEEN_HANDLE_AND_BAR_CENTERS = -20f
+private const val HANDLE_TRANSLATION = -30f
+private const val BAR_TRANSLATION_Y_FOR_TASKBAR = -50f
+private const val BAR_TRANSLATION_Y_FOR_HOTSEAT = -40f

quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt

@@ -18,19 +18,19 @@ package com.android.launcher3.taskbar.overlay
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.view.MotionEvent
+import androidx.test.annotation.UiThreadTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingView.TYPE_OPTIONS_POPUP
 import com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_ALL_APPS
 import com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_OVERLAY_PROXY
 import com.android.launcher3.AbstractFloatingView.hasOpenView
 import com.android.launcher3.taskbar.TaskbarActivityContext
-import com.android.launcher3.taskbar.TaskbarControllerTestUtil.runOnMainSync
-import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
-import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
-import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
+import com.android.launcher3.taskbar.TaskbarUnitTestRule
+import com.android.launcher3.taskbar.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.android.launcher3.util.TestUtil.getOnUiThread
+import com.android.launcher3.views.BaseDragLayer
 import com.android.systemui.shared.system.TaskStackChangeListeners
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
@@ -41,185 +41,193 @@ import org.junit.runner.RunWith
 @EmulatedDevices(["pixelFoldable2023"])
 class TaskbarOverlayControllerTest {
 
-    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
-    @get:Rule(order = 1) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule val taskbarUnitTestRule = TaskbarUnitTestRule()
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     private val taskbarContext: TaskbarActivityContext
         get() = taskbarUnitTestRule.activityContext
 
     @Test
+    @UiThreadTest
     fun testRequestWindow_twice_reusesWindow() {
-        val (context1, context2) =
-            getOnUiThread {
-                Pair(overlayController.requestWindow(), overlayController.requestWindow())
-            }
+        val context1 = overlayController.requestWindow()
+        val context2 = overlayController.requestWindow()
         assertThat(context1).isSameInstanceAs(context2)
     }
 
     @Test
+    @UiThreadTest
     fun testRequestWindow_afterHidingExistingWindow_createsNewWindow() {
-        val context1 = getOnUiThread { overlayController.requestWindow() }
-        runOnMainSync { overlayController.hideWindow() }
+        val context1 = overlayController.requestWindow()
+        overlayController.hideWindow()
 
-        val context2 = getOnUiThread { overlayController.requestWindow() }
+        val context2 = overlayController.requestWindow()
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
+    @UiThreadTest
     fun testRequestWindow_afterHidingOverlay_createsNewWindow() {
-        val context1 = getOnUiThread { overlayController.requestWindow() }
-        runOnMainSync {
-            TestOverlayView.show(context1)
-            overlayController.hideWindow()
-        }
+        val context1 = overlayController.requestWindow()
+        TestOverlayView.show(context1)
+        overlayController.hideWindow()
 
-        val context2 = getOnUiThread { overlayController.requestWindow() }
+        val context2 = overlayController.requestWindow()
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
+    @UiThreadTest
     fun testRequestWindow_addsProxyView() {
-        runOnMainSync { TestOverlayView.show(overlayController.requestWindow()) }
+        TestOverlayView.show(overlayController.requestWindow())
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
+    @UiThreadTest
     fun testRequestWindow_closeProxyView_closesOverlay() {
-        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
-        runOnMainSync {
-            AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
-        }
+        val overlay = TestOverlayView.show(overlayController.requestWindow())
+        AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testRequestWindow_attachesDragLayer() {
-        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
+        lateinit var dragLayer: BaseDragLayer<*>
+        getInstrumentation().runOnMainSync {
+            dragLayer = overlayController.requestWindow().dragLayer
+        }
+
         // Allow drag layer to attach before checking.
-        runOnMainSync { assertThat(dragLayer.isAttachedToWindow).isTrue() }
+        getInstrumentation().runOnMainSync { assertThat(dragLayer.isAttachedToWindow).isTrue() }
     }
 
     @Test
+    @UiThreadTest
     fun testHideWindow_closesOverlay() {
-        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
-        runOnMainSync { overlayController.hideWindow() }
+        val overlay = TestOverlayView.show(overlayController.requestWindow())
+        overlayController.hideWindow()
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testHideWindow_detachesDragLayer() {
-        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
+        lateinit var dragLayer: BaseDragLayer<*>
+        getInstrumentation().runOnMainSync {
+            dragLayer = overlayController.requestWindow().dragLayer
+        }
 
         // Wait for drag layer to be attached to window before hiding.
-        runOnMainSync {
+        getInstrumentation().runOnMainSync {
             overlayController.hideWindow()
             assertThat(dragLayer.isAttachedToWindow).isFalse()
         }
     }
 
     @Test
+    @UiThreadTest
     fun testTwoOverlays_closeOne_windowStaysOpen() {
-        val (overlay1, overlay2) =
-            getOnUiThread {
-                val context = overlayController.requestWindow()
-                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
-            }
+        val context = overlayController.requestWindow()
+        val overlay1 = TestOverlayView.show(context)
+        val overlay2 = TestOverlayView.show(context)
 
-        runOnMainSync { overlay1.close(false) }
+        overlay1.close(false)
         assertThat(overlay2.isOpen).isTrue()
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
+    @UiThreadTest
     fun testTwoOverlays_closeAll_closesWindow() {
-        val (overlay1, overlay2) =
-            getOnUiThread {
-                val context = overlayController.requestWindow()
-                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
-            }
+        val context = overlayController.requestWindow()
+        val overlay1 = TestOverlayView.show(context)
+        val overlay2 = TestOverlayView.show(context)
 
-        runOnMainSync {
-            overlay1.close(false)
-            overlay2.close(false)
-        }
+        overlay1.close(false)
+        overlay2.close(false)
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isFalse()
     }
 
     @Test
+    @UiThreadTest
     fun testRecreateTaskbar_closesWindow() {
-        runOnMainSync { TestOverlayView.show(overlayController.requestWindow()) }
+        TestOverlayView.show(overlayController.requestWindow())
         taskbarUnitTestRule.recreateTaskbar()
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isFalse()
     }
 
     @Test
     fun testTaskMovedToFront_closesOverlay() {
-        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        lateinit var overlay: TestOverlayView
+        getInstrumentation().runOnMainSync {
+            overlay = TestOverlayView.show(overlayController.requestWindow())
+        }
+
         TaskStackChangeListeners.getInstance().listenerImpl.onTaskMovedToFront(RunningTaskInfo())
         // Make sure TaskStackChangeListeners' Handler posts the callback before checking state.
-        runOnMainSync { assertThat(overlay.isOpen).isFalse() }
+        getInstrumentation().runOnMainSync { assertThat(overlay.isOpen).isFalse() }
     }
 
     @Test
     fun testTaskStackChanged_allAppsClosed_overlayStaysOpen() {
-        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
-        runOnMainSync { taskbarContext.controllers.sharedState?.allAppsVisible = false }
+        lateinit var overlay: TestOverlayView
+        getInstrumentation().runOnMainSync {
+            overlay = TestOverlayView.show(overlayController.requestWindow())
+            taskbarContext.controllers.sharedState?.allAppsVisible = false
+        }
 
         TaskStackChangeListeners.getInstance().listenerImpl.onTaskStackChanged()
-        runOnMainSync { assertThat(overlay.isOpen).isTrue() }
+        getInstrumentation().runOnMainSync { assertThat(overlay.isOpen).isTrue() }
     }
 
     @Test
     fun testTaskStackChanged_allAppsOpen_closesOverlay() {
-        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
-        runOnMainSync { taskbarContext.controllers.sharedState?.allAppsVisible = true }
+        lateinit var overlay: TestOverlayView
+        getInstrumentation().runOnMainSync {
+            overlay = TestOverlayView.show(overlayController.requestWindow())
+            taskbarContext.controllers.sharedState?.allAppsVisible = true
+        }
 
         TaskStackChangeListeners.getInstance().listenerImpl.onTaskStackChanged()
-        runOnMainSync { assertThat(overlay.isOpen).isFalse() }
+        getInstrumentation().runOnMainSync { assertThat(overlay.isOpen).isFalse() }
     }
 
     @Test
+    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayNotRebindSafe_closesOverlay() {
-        val context = getOnUiThread { overlayController.requestWindow() }
-        val overlay = getOnUiThread {
-            TestOverlayView.show(context).apply { type = TYPE_OPTIONS_POPUP }
-        }
+        val overlayContext = overlayController.requestWindow()
+        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_OPTIONS_POPUP }
 
-        runOnMainSync {
-            overlayController.updateLauncherDeviceProfile(
-                overlayController.launcherDeviceProfile
-                    .toBuilder(context)
-                    .setGestureMode(false)
-                    .build()
-            )
-        }
+        overlayController.updateLauncherDeviceProfile(
+            overlayController.launcherDeviceProfile
+                .toBuilder(overlayContext)
+                .setGestureMode(false)
+                .build()
+        )
 
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
+    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayRebindSafe_overlayStaysOpen() {
-        val context = getOnUiThread { overlayController.requestWindow() }
-        val overlay = getOnUiThread {
-            TestOverlayView.show(context).apply { type = TYPE_TASKBAR_ALL_APPS }
-        }
+        val overlayContext = overlayController.requestWindow()
+        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_TASKBAR_ALL_APPS }
 
-        runOnMainSync {
-            overlayController.updateLauncherDeviceProfile(
-                overlayController.launcherDeviceProfile
-                    .toBuilder(context)
-                    .setGestureMode(false)
-                    .build()
-            )
-        }
+        overlayController.updateLauncherDeviceProfile(
+            overlayController.launcherDeviceProfile
+                .toBuilder(overlayContext)
+                .setGestureMode(false)
+                .build()
+        )
 
         assertThat(overlay.isOpen).isTrue()
     }
 
     private class TestOverlayView
-    private constructor(private val overlayContext: TaskbarOverlayContext) :
-        AbstractFloatingView(overlayContext, null) {
+    private constructor(
+        private val overlayContext: TaskbarOverlayContext,
+    ) : AbstractFloatingView(overlayContext, null) {
 
         var type = TYPE_OPTIONS_POPUP
 

quickstep/tests/multivalentTests/src/com/android/quickstep/AllAppsActionManagerTest.kt

@@ -18,59 +18,32 @@ package com.android.quickstep
 
 import android.app.PendingIntent
 import android.content.IIntentSender
-import android.provider.Settings
-import android.provider.Settings.Secure.USER_SETUP_COMPLETE
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.launcher3.dagger.LauncherAppComponent
-import com.android.launcher3.dagger.LauncherAppSingleton
-import com.android.launcher3.util.AllModulesForTest
+import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
-import com.android.launcher3.util.SandboxApplication
-import com.android.launcher3.util.SettingsCache
-import com.android.launcher3.util.SettingsCacheSandbox
 import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
-import dagger.BindsInstance
-import dagger.Component
 import java.util.concurrent.Semaphore
 import java.util.concurrent.TimeUnit.SECONDS
-import org.junit.After
-import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 private const val TIMEOUT = 5L
-private val USER_SETUP_COMPLETE_URI = Settings.Secure.getUriFor(USER_SETUP_COMPLETE)
 
 @RunWith(AndroidJUnit4::class)
 class AllAppsActionManagerTest {
     private val callbackSemaphore = Semaphore(0)
     private val bgExecutor = UI_HELPER_EXECUTOR
 
-    @get:Rule val context = SandboxApplication()
-
-    private val settingsCacheSandbox =
-        SettingsCacheSandbox().also { it[USER_SETUP_COMPLETE_URI] = 1 }
-
-    private val allAppsActionManager by
-        lazy(LazyThreadSafetyMode.NONE) {
-            AllAppsActionManager(context, bgExecutor) {
-                callbackSemaphore.release()
-                PendingIntent(IIntentSender.Default())
-            }
+    private val allAppsActionManager =
+        AllAppsActionManager(
+            InstrumentationRegistry.getInstrumentation().targetContext,
+            bgExecutor,
+        ) {
+            callbackSemaphore.release()
+            PendingIntent(IIntentSender.Default())
         }
 
-    @Before
-    fun initDaggerComponent() {
-        context.initDaggerComponent(
-            DaggerAllAppsActionManagerTestComponent.builder()
-                .bindSettingsCache(settingsCacheSandbox.cache)
-        )
-    }
-
-    @After fun destroyManager() = allAppsActionManager.onDestroy()
-
     @Test
     fun taskbarPresent_actionRegistered() {
         allAppsActionManager.isTaskbarPresent = true
@@ -115,50 +88,4 @@ class AllAppsActionManagerTest {
         assertThat(callbackSemaphore.tryAcquire(TIMEOUT, SECONDS)).isTrue()
         assertThat(allAppsActionManager.isActionRegistered).isTrue()
     }
-
-    @Test
-    fun taskbarPresent_userSetupIncomplete_actionUnregistered() {
-        settingsCacheSandbox[USER_SETUP_COMPLETE_URI] = 0
-        allAppsActionManager.isTaskbarPresent = true
-        assertThat(allAppsActionManager.isActionRegistered).isFalse()
-    }
-
-    @Test
-    fun taskbarPresent_setupUiVisible_actionUnregistered() {
-        allAppsActionManager.isSetupUiVisible = true
-        allAppsActionManager.isTaskbarPresent = true
-        assertThat(allAppsActionManager.isActionRegistered).isFalse()
-    }
-
-    @Test
-    fun taskbarPresent_userSetupCompleted_actionRegistered() {
-        settingsCacheSandbox[USER_SETUP_COMPLETE_URI] = 0
-        allAppsActionManager.isTaskbarPresent = true
-
-        settingsCacheSandbox[USER_SETUP_COMPLETE_URI] = 1
-        assertThat(callbackSemaphore.tryAcquire(TIMEOUT, SECONDS)).isTrue()
-        assertThat(allAppsActionManager.isActionRegistered).isTrue()
-    }
-
-    @Test
-    fun taskbarPresent_setupUiDismissed_actionRegistered() {
-        allAppsActionManager.isSetupUiVisible = true
-        allAppsActionManager.isTaskbarPresent = true
-
-        allAppsActionManager.isSetupUiVisible = false
-        assertThat(callbackSemaphore.tryAcquire(TIMEOUT, SECONDS)).isTrue()
-        assertThat(allAppsActionManager.isActionRegistered).isTrue()
-    }
-}
-
-@LauncherAppSingleton
-@Component(modules = [AllModulesForTest::class])
-interface AllAppsActionManagerTestComponent : LauncherAppComponent {
-
-    @Component.Builder
-    interface Builder : LauncherAppComponent.Builder {
-        @BindsInstance fun bindSettingsCache(settingsCache: SettingsCache): Builder
-
-        override fun build(): AllAppsActionManagerTestComponent
-    }
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/FullscreenDrawParamsTest.kt

@@ -20,17 +20,21 @@ import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.FakeInvariantDeviceProfileTest
 import com.android.quickstep.util.TaskCornerRadius
+import com.android.quickstep.views.TaskView.FullscreenDrawParams
 import com.android.systemui.shared.system.QuickStepContract
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.kotlin.mock
+import org.mockito.Mockito.doReturn
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.spy
 
-/** Test for [FullscreenDrawParams] class. */
+/** Test for FullscreenDrawParams class. */
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class FullscreenDrawParamsTest : FakeInvariantDeviceProfileTest() {
+
     private lateinit var params: FullscreenDrawParams
 
     @Before
@@ -42,108 +46,115 @@ class FullscreenDrawParamsTest : FakeInvariantDeviceProfileTest() {
     fun setStartProgress_correctCornerRadiusForTablet() {
         initializeVarsForTablet()
 
-        params.setProgress(fullscreenProgress = 0f, parentScale = 1.0f, taskViewScale = 1.0f)
+        params.setProgress(
+            /* fullscreenProgress= */ 0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
 
         val expectedRadius = TaskCornerRadius.get(context)
-        assertThat(params.currentCornerRadius).isEqualTo(expectedRadius)
+        assertThat(params.currentDrawnCornerRadius).isEqualTo(expectedRadius)
     }
 
     @Test
     fun setFullProgress_correctCornerRadiusForTablet() {
         initializeVarsForTablet()
 
-        params.setProgress(fullscreenProgress = 1.0f, parentScale = 1f, taskViewScale = 1f)
+        params.setProgress(
+            /* fullscreenProgress= */ 1.0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
 
         val expectedRadius = QuickStepContract.getWindowCornerRadius(context)
-        assertThat(params.currentCornerRadius).isEqualTo(expectedRadius)
+        assertThat(params.currentDrawnCornerRadius).isEqualTo(expectedRadius)
     }
 
     @Test
     fun setStartProgress_correctCornerRadiusForPhone() {
         initializeVarsForPhone()
 
-        params.setProgress(fullscreenProgress = 0f, parentScale = 1f, taskViewScale = 1f)
+        params.setProgress(
+            /* fullscreenProgress= */ 0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
 
         val expectedRadius = TaskCornerRadius.get(context)
-        assertThat(params.currentCornerRadius).isEqualTo(expectedRadius)
+        assertThat(params.currentDrawnCornerRadius).isEqualTo(expectedRadius)
     }
 
     @Test
     fun setFullProgress_correctCornerRadiusForPhone() {
         initializeVarsForPhone()
 
-        params.setProgress(fullscreenProgress = 1.0f, parentScale = 1f, taskViewScale = 1f)
+        params.setProgress(
+            /* fullscreenProgress= */ 1.0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
 
         val expectedRadius = QuickStepContract.getWindowCornerRadius(context)
-        assertThat(params.currentCornerRadius).isEqualTo(expectedRadius)
+        assertThat(params.currentDrawnCornerRadius).isEqualTo(expectedRadius)
     }
 
     @Test
     fun setStartProgress_correctCornerRadiusForMultiDisplay() {
-        val display1Context = mock<Context>()
-        val display2Context = mock<Context>()
-        val display1TaskRadius = TASK_CORNER_RADIUS + 1
-        val display2TaskRadius = TASK_CORNER_RADIUS + 2
+        val display1Context = context
+        val display2Context = mock(Context::class.java)
+        val spyParams = spy(params)
 
-        val params =
-            FullscreenDrawParams(
-                context,
-                taskCornerRadiusProvider = { context ->
-                    when (context) {
-                        display1Context -> display1TaskRadius
-                        display2Context -> display2TaskRadius
-                        else -> TASK_CORNER_RADIUS
-                    }
-                },
-                windowCornerRadiusProvider = { 0f },
-            )
+        val display1TaskRadius = TaskCornerRadius.get(display1Context)
+        val display1WindowRadius = QuickStepContract.getWindowCornerRadius(display1Context)
+        val display2TaskRadius = display1TaskRadius * 2 + 1 // Arbitrarily different.
+        val display2WindowRadius = display1WindowRadius * 2 + 1 // Arbitrarily different.
+        doReturn(display2TaskRadius).`when`(spyParams).computeTaskCornerRadius(display2Context)
+        doReturn(display2WindowRadius).`when`(spyParams).computeWindowCornerRadius(display2Context)
 
-        params.setProgress(fullscreenProgress = 0f, parentScale = 1f, taskViewScale = 1f)
-        assertThat(params.currentCornerRadius).isEqualTo(TASK_CORNER_RADIUS)
+        spyParams.updateCornerRadius(display1Context)
+        spyParams.setProgress(
+            /* fullscreenProgress= */ 0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
+        assertThat(spyParams.currentDrawnCornerRadius).isEqualTo(display1TaskRadius)
 
-        params.updateCornerRadius(display1Context)
-        params.setProgress(fullscreenProgress = 0f, parentScale = 1f, taskViewScale = 1f)
-        assertThat(params.currentCornerRadius).isEqualTo(display1TaskRadius)
-
-        params.updateCornerRadius(display2Context)
-        params.setProgress(fullscreenProgress = 0f, parentScale = 1f, taskViewScale = 1f)
-        assertThat(params.currentCornerRadius).isEqualTo(display2TaskRadius)
+        spyParams.updateCornerRadius(display2Context)
+        spyParams.setProgress(
+            /* fullscreenProgress= */ 0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
+        assertThat(spyParams.currentDrawnCornerRadius).isEqualTo(display2TaskRadius)
     }
 
     @Test
     fun setFullProgress_correctCornerRadiusForMultiDisplay() {
-        val display1Context = mock<Context>()
-        val display2Context = mock<Context>()
-        val display1WindowRadius = WINDOW_CORNER_RADIUS + 1
-        val display2WindowRadius = WINDOW_CORNER_RADIUS + 2
+        val display1Context = context
+        val display2Context = mock(Context::class.java)
+        val spyParams = spy(params)
 
-        val params =
-            FullscreenDrawParams(
-                context,
-                taskCornerRadiusProvider = { 0f },
-                windowCornerRadiusProvider = { context ->
-                    when (context) {
-                        display1Context -> display1WindowRadius
-                        display2Context -> display2WindowRadius
-                        else -> WINDOW_CORNER_RADIUS
-                    }
-                },
-            )
+        val display1TaskRadius = TaskCornerRadius.get(display1Context)
+        val display1WindowRadius = QuickStepContract.getWindowCornerRadius(display1Context)
+        val display2TaskRadius = display1TaskRadius * 2 + 1 // Arbitrarily different.
+        val display2WindowRadius = display1WindowRadius * 2 + 1 // Arbitrarily different.
+        doReturn(display2TaskRadius).`when`(spyParams).computeTaskCornerRadius(display2Context)
+        doReturn(display2WindowRadius).`when`(spyParams).computeWindowCornerRadius(display2Context)
 
-        params.setProgress(fullscreenProgress = 1f, parentScale = 1f, taskViewScale = 1f)
-        assertThat(params.currentCornerRadius).isEqualTo(WINDOW_CORNER_RADIUS)
+        spyParams.updateCornerRadius(display1Context)
+        spyParams.setProgress(
+            /* fullscreenProgress= */ 1.0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f
+        )
+        assertThat(spyParams.currentDrawnCornerRadius).isEqualTo(display1WindowRadius)
 
-        params.updateCornerRadius(display1Context)
-        params.setProgress(fullscreenProgress = 1f, parentScale = 1f, taskViewScale = 1f)
-        assertThat(params.currentCornerRadius).isEqualTo(display1WindowRadius)
-
-        params.updateCornerRadius(display2Context)
-        params.setProgress(fullscreenProgress = 1f, parentScale = 1f, taskViewScale = 1f)
-        assertThat(params.currentCornerRadius).isEqualTo(display2WindowRadius)
-    }
-
-    companion object {
-        const val TASK_CORNER_RADIUS = 56f
-        const val WINDOW_CORNER_RADIUS = 32f
+        spyParams.updateCornerRadius(display2Context)
+        spyParams.setProgress(
+            /* fullscreenProgress= */ 1.0f,
+            /* parentScale= */ 1.0f,
+            /* taskViewScale= */ 1.0f,
+        )
+        assertThat(spyParams.currentDrawnCornerRadius).isEqualTo(display2WindowRadius)
     }
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/orientation/LandscapePagedViewHandlerTest.kt

@@ -43,12 +43,12 @@ class LandscapePagedViewHandlerTest {
         if (isEnabled) {
             setFlagsRule.enableFlags(
                 Flags.FLAG_ENABLE_GRID_ONLY_OVERVIEW,
-                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU,
+                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU
             )
         } else {
             setFlagsRule.disableFlags(
                 Flags.FLAG_ENABLE_GRID_ONLY_OVERVIEW,
-                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU,
+                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU
             )
         }
     }
@@ -62,7 +62,6 @@ class LandscapePagedViewHandlerTest {
             isRTL,
             OVERVIEW_TASK_MARGIN_PX,
             DIVIDER_SIZE_PX,
-            oneIconHiddenDueToSmallWidth = false,
         )
     }
 
@@ -108,8 +107,14 @@ class LandscapePagedViewHandlerTest {
 
         val (topLeftY, bottomRightY) = getSplitIconsPosition(isRTL = true)
 
-        assertThat(topLeftY).isEqualTo(-316)
-        assertThat(bottomRightY).isEqualTo(0)
+        // TODO(b/326377497): When started in fake seascape and rotated to landscape,
+        //  the icon chips are in RTL and wrongly positioned at the right side of the snapshot.
+        //  Top-Left app chip should be placed at the top left of the first snapshot, but because
+        //  this issue, it's displayed at the top-right of the second snapshot.
+        //  The Bottom-Right app chip is displayed at the top-right of the first snapshot because
+        //  of this issue.
+        assertThat(topLeftY).isEqualTo(0)
+        assertThat(bottomRightY).isEqualTo(-316)
     }
 
     /** Test updateSplitIconsPosition */

quickstep/tests/multivalentTests/src/com/android/quickstep/orientation/SeascapePagedViewHandlerTest.kt

@@ -43,12 +43,12 @@ class SeascapePagedViewHandlerTest {
         if (isEnabled) {
             setFlagsRule.enableFlags(
                 Flags.FLAG_ENABLE_GRID_ONLY_OVERVIEW,
-                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU,
+                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU
             )
         } else {
             setFlagsRule.disableFlags(
                 Flags.FLAG_ENABLE_GRID_ONLY_OVERVIEW,
-                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU,
+                Flags.FLAG_ENABLE_OVERVIEW_ICON_MENU
             )
         }
     }
@@ -62,7 +62,6 @@ class SeascapePagedViewHandlerTest {
             isRTL,
             OVERVIEW_TASK_MARGIN_PX,
             DIVIDER_SIZE_PX,
-            oneIconHiddenDueToSmallWidth = false,
         )
     }
 
@@ -110,6 +109,12 @@ class SeascapePagedViewHandlerTest {
 
         val (topLeftY, bottomRightY) = getSplitIconsPosition(isRTL = true)
 
+        // TODO(b/326377497): When started in fake seascape and rotated to landscape,
+        //  the icon chips are in RTL and wrongly positioned at the right side of the snapshot.
+        //  Top-Left app chip should be placed at the top left of the first snapshot, but because
+        //  this issue, it's displayed at the top-right of the second snapshot.
+        //  The Bottom-Right app chip is displayed at the top-right of the first snapshot because
+        //  of this issue.
         assertThat(topLeftY).isEqualTo(316)
         assertThat(bottomRightY).isEqualTo(0)
     }
@@ -161,7 +166,7 @@ class SeascapePagedViewHandlerTest {
         `when`(iconView.layoutParams).thenReturn(frameLayout)
 
         sut.updateSplitIconsPosition(iconView, expectedTranslationY, false)
-        assertThat(frameLayout.gravity).isEqualTo(Gravity.BOTTOM or Gravity.END)
+        assertThat(frameLayout.gravity).isEqualTo(Gravity.BOTTOM or Gravity.START)
         verify(iconView).setSplitTranslationX(0f)
         verify(iconView).setSplitTranslationY(expectedTranslationY.toFloat())
     }
@@ -176,7 +181,7 @@ class SeascapePagedViewHandlerTest {
         `when`(iconView.layoutParams).thenReturn(frameLayout)
 
         sut.updateSplitIconsPosition(iconView, expectedTranslationY, true)
-        assertThat(frameLayout.gravity).isEqualTo(Gravity.TOP or Gravity.START)
+        assertThat(frameLayout.gravity).isEqualTo(Gravity.TOP or Gravity.END)
         verify(iconView).setSplitTranslationX(0f)
         verify(iconView).setSplitTranslationY(expectedTranslationY.toFloat())
     }

quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeRecentTasksDataSource.kt

@@ -20,12 +20,10 @@ import com.android.quickstep.util.GroupTask
 import java.util.function.Consumer
 
 class FakeRecentTasksDataSource : RecentTasksDataSource {
-    private var taskList: List<GroupTask> = listOf()
+    var taskList: List<GroupTask> = listOf()
 
     override fun getTasks(callback: Consumer<List<GroupTask>>?): Int {
-        // Makes a copy of the GroupTask to create a new GroupTask instance and to simulate
-        // RecentsModel::getTasks behavior.
-        callback?.accept(taskList.map { it.copy() })
+        callback?.accept(taskList)
         return 0
     }
 

quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt

@@ -17,41 +17,36 @@
 package com.android.quickstep.recents.data
 
 import android.graphics.Bitmap
+import com.android.launcher3.util.CancellableTask
 import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
-import kotlinx.coroutines.yield
+import java.util.function.Consumer
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 class FakeTaskThumbnailDataSource : TaskThumbnailDataSource {
 
-    val taskIdToBitmap: MutableMap<Int, Bitmap> =
-        (0..10).associateWith { mock<Bitmap>() }.toMutableMap()
-    private val completionPrevented: MutableSet<Int> = mutableSetOf()
-    private val getThumbnailCalls = mutableMapOf<Int, Int>()
-
-    var highResEnabled = true
+    val taskIdToBitmap: Map<Int, Bitmap> = (0..10).associateWith { mock() }
+    val taskIdToUpdatingTask: MutableMap<Int, () -> Unit> = mutableMapOf()
+    var shouldLoadSynchronously: Boolean = true
 
     /** Retrieves and sets a thumbnail on [task] from [taskIdToBitmap]. */
-    override suspend fun getThumbnail(task: Task): ThumbnailData {
-        getThumbnailCalls[task.key.id] = (getThumbnailCalls[task.key.id] ?: 0) + 1
-
-        while (task.key.id in completionPrevented) {
-            yield()
+    override fun updateThumbnailInBackground(
+        task: Task,
+        callback: Consumer<ThumbnailData>
+    ): CancellableTask<ThumbnailData>? {
+        val thumbnailData = mock<ThumbnailData>()
+        whenever(thumbnailData.thumbnail).thenReturn(taskIdToBitmap[task.key.id])
+        val wrappedCallback = {
+            task.thumbnail = thumbnailData
+            callback.accept(thumbnailData)
         }
-        return ThumbnailData(
-            thumbnail = taskIdToBitmap[task.key.id],
-            reducedResolution = !highResEnabled,
-        )
-    }
-
-    fun getNumberOfGetThumbnailCalls(taskId: Int): Int = getThumbnailCalls[taskId] ?: 0
-
-    fun preventThumbnailLoad(taskId: Int) {
-        completionPrevented.add(taskId)
-    }
-
-    fun completeLoadingForTask(taskId: Int) {
-        completionPrevented.remove(taskId)
+        if (shouldLoadSynchronously) {
+            wrappedCallback()
+        } else {
+            taskIdToUpdatingTask[task.key.id] = wrappedCallback
+        }
+        return null
     }
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt

@@ -16,55 +16,25 @@
 
 package com.android.quickstep.recents.data
 
-import android.graphics.drawable.Drawable
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 
 class FakeTasksRepository : RecentTasksRepository {
     private var thumbnailDataMap: Map<Int, ThumbnailData> = emptyMap()
-    private var taskIconDataMap: Map<Int, FakeIconData> = emptyMap()
     private var tasks: MutableStateFlow<List<Task>> = MutableStateFlow(emptyList())
-    private var visibleTasks: MutableStateFlow<Set<Int>> = MutableStateFlow(emptySet())
+    private var visibleTasks: MutableStateFlow<List<Int>> = MutableStateFlow(emptyList())
 
     override fun getAllTaskData(forceRefresh: Boolean): Flow<List<Task>> = tasks
 
     override fun getTaskDataById(taskId: Int): Flow<Task?> =
-        combine(getAllTaskData(), visibleTasks) { taskList, visibleTasks ->
-                taskList.filter { visibleTasks.contains(it.key.id) }
-            }
-            .map { taskList ->
-                val task = taskList.firstOrNull { it.key.id == taskId } ?: return@map null
-                Task(task).apply {
-                    thumbnail = task.thumbnail
-                    icon = task.icon
-                    titleDescription = task.titleDescription
-                    title = task.title
-                }
-            }
+        getAllTaskData().map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
 
-    override fun getThumbnailById(taskId: Int): Flow<ThumbnailData?> =
-        getTaskDataById(taskId).map { it?.thumbnail }
-
-    override fun getCurrentThumbnailById(taskId: Int): ThumbnailData? =
-        tasks.value.firstOrNull { it.key.id == taskId }?.thumbnail
-
-    override fun setVisibleTasks(visibleTaskIdList: Set<Int>) {
+    override fun setVisibleTasks(visibleTaskIdList: List<Int>) {
         visibleTasks.value = visibleTaskIdList
-        tasks.value =
-            tasks.value.map {
-                it.apply {
-                    thumbnail = thumbnailDataMap[it.key.id]
-                    taskIconDataMap[it.key.id]?.let { data ->
-                        title = data.title
-                        titleDescription = data.titleDescription
-                        icon = data.icon
-                    }
-                }
-            }
+        tasks.value = tasks.value.map { it.apply { thumbnail = thumbnailDataMap[it.key.id] } }
     }
 
     fun seedTasks(tasks: List<Task>) {
@@ -74,15 +44,4 @@ class FakeTasksRepository : RecentTasksRepository {
     fun seedThumbnailData(thumbnailDataMap: Map<Int, ThumbnailData>) {
         this.thumbnailDataMap = thumbnailDataMap
     }
-
-    fun seedIconData(id: Int, title: String, contentDescription: String, icon: Drawable) {
-        val iconData = FakeIconData(icon, contentDescription, title)
-        this.taskIconDataMap = mapOf(id to iconData)
-    }
-
-    private data class FakeIconData(
-        val icon: Drawable,
-        val titleDescription: String,
-        val title: String,
-    )
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt

@@ -18,455 +18,133 @@ package com.android.quickstep.recents.data
 
 import android.content.ComponentName
 import android.content.Intent
-import android.graphics.Bitmap
-import android.graphics.Rect
-import android.graphics.drawable.Drawable
-import android.view.Display.DEFAULT_DISPLAY
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.launcher3.util.SplitConfigurationOptions
-import com.android.launcher3.util.TestDispatcherProvider
+import com.android.quickstep.TaskIconCache
 import com.android.quickstep.util.DesktopTask
-import com.android.quickstep.util.SingleTask
-import com.android.quickstep.util.SplitTask
+import com.android.quickstep.util.GroupTask
 import com.android.systemui.shared.recents.model.Task
-import com.android.systemui.shared.recents.model.ThumbnailData
-import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.drop
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
-import org.junit.runner.RunWith
-import org.mockito.Mockito.spy
 import org.mockito.kotlin.mock
-import org.mockito.kotlin.verify
-import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
-@RunWith(AndroidJUnit4::class)
 class TasksRepositoryTest {
     private val tasks = (0..5).map(::createTaskWithId)
     private val defaultTaskList =
         listOf(
-            SingleTask(tasks[0]),
-            SplitTask(
-                tasks[1],
-                tasks[2],
-                SplitConfigurationOptions.SplitBounds(
-                    /* leftTopBounds = */ Rect(),
-                    /* rightBottomBounds = */ Rect(),
-                    /* leftTopTaskId = */ -1,
-                    /* rightBottomTaskId = */ -1,
-                    /* snapPosition = */ SNAP_TO_2_50_50,
-                ),
-            ),
-            DesktopTask(deskId = 0, DEFAULT_DISPLAY, tasks.subList(3, 6)),
+            GroupTask(tasks[0]),
+            GroupTask(tasks[1], tasks[2], null),
+            DesktopTask(tasks.subList(3, 6))
         )
     private val recentsModel = FakeRecentTasksDataSource()
     private val taskThumbnailDataSource = FakeTaskThumbnailDataSource()
-    private val taskIconDataSource = FakeTaskIconDataSource()
-    private val taskVisualsChangeNotifier = FakeTaskVisualsChangeNotifier()
-    private val highResLoadingStateNotifier = FakeHighResLoadingStateNotifier()
-    private val taskVisualsChangedDelegate =
-        spy(TaskVisualsChangedDelegateImpl(taskVisualsChangeNotifier, highResLoadingStateNotifier))
+    private val taskIconCache = mock<TaskIconCache>()
 
-    private val dispatcher = UnconfinedTestDispatcher()
-    private val testScope = TestScope(dispatcher)
     private val systemUnderTest =
-        TasksRepository(
-            recentsModel,
-            taskThumbnailDataSource,
-            taskIconDataSource,
-            taskVisualsChangedDelegate,
-            testScope.backgroundScope,
-            TestDispatcherProvider(dispatcher),
-        )
+        TasksRepository(recentsModel, taskThumbnailDataSource, taskIconCache)
 
     @Test
-    fun getAllTaskDataReturnsFlattenedListOfTasks() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            assertThat(systemUnderTest.getAllTaskData(forceRefresh = true).first()).isEqualTo(tasks)
-        }
+    fun getAllTaskDataReturnsFlattenedListOfTasks() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+
+        assertThat(systemUnderTest.getAllTaskData(forceRefresh = true).first()).isEqualTo(tasks)
+    }
 
     @Test
-    fun getTaskDataByIdReturnsSpecificTask() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun getTaskDataByIdReturnsSpecificTask() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        systemUnderTest.getAllTaskData(forceRefresh = true)
 
-            assertThat(systemUnderTest.getTaskDataById(2).first()).isEqualTo(tasks[2])
-        }
+        assertThat(systemUnderTest.getTaskDataById(2).first()).isEqualTo(tasks[2])
+    }
 
     @Test
-    fun getThumbnailByIdReturnsNullWithNoLoadedThumbnails() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun setVisibleTasksPopulatesThumbnails() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        systemUnderTest.getAllTaskData(forceRefresh = true)
 
-            assertThat(systemUnderTest.getThumbnailById(1).first()).isNull()
-        }
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getTaskDataById(1).drop(1).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap1)
+        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+    }
 
     @Test
-    fun getCurrentThumbnailByIdReturnsNullWithNoLoadedThumbnails() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun changingVisibleTasksContainsAlreadyPopulatedThumbnails() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        systemUnderTest.getAllTaskData(forceRefresh = true)
 
-            assertThat(systemUnderTest.getCurrentThumbnailById(1)).isNull()
-        }
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getTaskDataById(2).drop(1).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+
+        // Prevent new loading of Bitmaps
+        taskThumbnailDataSource.shouldLoadSynchronously = false
+        systemUnderTest.setVisibleTasks(listOf(2, 3))
+
+        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+    }
 
     @Test
-    fun getThumbnailByIdReturnsThumbnailWithLoadedThumbnails() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+    fun retrievedThumbnailsAreDiscardedWhenTaskBecomesInvisible() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        systemUnderTest.getAllTaskData(forceRefresh = true)
 
-            systemUnderTest.setVisibleTasks(setOf(1))
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
 
-            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
-        }
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getTaskDataById(2).drop(1).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+
+        // Prevent new loading of Bitmaps
+        taskThumbnailDataSource.shouldLoadSynchronously = false
+        systemUnderTest.setVisibleTasks(listOf(0, 1))
+
+        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail).isNull()
+    }
 
     @Test
-    fun whenThumbnailIsLoaded_getAllTaskData_usesPreviousLoadedThumbnailAndIcon() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+    fun retrievedThumbnailsCauseEmissionOnTaskDataFlow() = runTest {
+        // Setup fakes
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        taskThumbnailDataSource.shouldLoadSynchronously = false
 
-            systemUnderTest.setVisibleTasks(setOf(1))
-            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
+        // Setup TasksRepository
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
 
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
+        // Assert there is no bitmap in first emission
+        val taskFlow = systemUnderTest.getTaskDataById(2)
+        val taskFlowValuesList = mutableListOf<Task?>()
+        backgroundScope.launch(UnconfinedTestDispatcher(testScheduler)) {
+            taskFlow.toList(taskFlowValuesList)
         }
+        assertThat(taskFlowValuesList[0]!!.thumbnail).isNull()
 
-    @Test
-    fun getAllTaskData_clearsPreviouslyLoadedImagesForRemovedTasks() =
-        testScope.runTest {
-            // Setup data
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        // Simulate bitmap loading after first emission
+        taskThumbnailDataSource.taskIdToUpdatingTask.getValue(2).invoke()
 
-            // Load images for task 1
-            systemUnderTest.setVisibleTasks(setOf(1))
-            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
-
-            // Remove task 1 from "all data"
-            recentsModel.seedTasks(
-                defaultTaskList.filterNot { groupTask -> groupTask.tasks.any { it.key.id == 1 } }
-            )
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            // Assert task 1 was fully removed
-            assertThat(systemUnderTest.getThumbnailById(1).first()?.thumbnail).isNull()
-            verify(taskVisualsChangedDelegate).unregisterTaskThumbnailChangedCallback(tasks[1].key)
-        }
-
-    @Test
-    fun getCurrentThumbnailByIdReturnsThumbnailWithLoadedThumbnails() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-
-            systemUnderTest.setVisibleTasks(setOf(1))
-
-            assertThat(systemUnderTest.getCurrentThumbnailById(1)?.thumbnail).isEqualTo(bitmap1)
-        }
-
-    @Test
-    fun setVisibleTasksPopulatesThumbnails() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
-
-            assertThat(systemUnderTest.getTaskDataById(1).first()!!.thumbnail!!.thumbnail)
-                .isEqualTo(bitmap1)
-            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
-                .isEqualTo(bitmap2)
-        }
-
-    @Test
-    fun setVisibleTasksPopulatesIcons() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
-
-            systemUnderTest
-                .getTaskDataById(1)
-                .first()!!
-                .assertHasIconDataFromSource(taskIconDataSource)
-            systemUnderTest
-                .getTaskDataById(2)
-                .first()!!
-                .assertHasIconDataFromSource(taskIconDataSource)
-        }
-
-    @Test
-    fun changingVisibleTasksContainsAlreadyPopulatedThumbnails() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
-
-            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
-                .isEqualTo(bitmap2)
-
-            // Prevent new loading of Bitmaps
-            taskThumbnailDataSource.preventThumbnailLoad(2)
-            systemUnderTest.setVisibleTasks(setOf(2, 3))
-
-            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
-                .isEqualTo(bitmap2)
-        }
-
-    @Test
-    fun changingVisibleTasksContainsAlreadyPopulatedIcons() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
-
-            systemUnderTest
-                .getTaskDataById(2)
-                .first()!!
-                .assertHasIconDataFromSource(taskIconDataSource)
-
-            // Prevent new loading of Drawables
-            taskIconDataSource.preventIconLoad(2)
-            systemUnderTest.setVisibleTasks(setOf(2, 3))
-
-            systemUnderTest
-                .getTaskDataById(2)
-                .first()!!
-                .assertHasIconDataFromSource(taskIconDataSource)
-        }
-
-    @Test
-    fun retrievedImagesAreDiscardedWhenTaskBecomesInvisible() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
-
-            val task2 = systemUnderTest.getTaskDataById(2).first()!!
-            assertThat(task2.thumbnail!!.thumbnail).isEqualTo(bitmap2)
-            task2.assertHasIconDataFromSource(taskIconDataSource)
-
-            systemUnderTest.setVisibleTasks(setOf(0, 1))
-
-            val task2AfterVisibleTasksChanged = systemUnderTest.getTaskDataById(2).first()!!
-            assertThat(task2AfterVisibleTasksChanged.thumbnail).isNull()
-            assertThat(task2AfterVisibleTasksChanged.icon).isNull()
-            assertThat(task2AfterVisibleTasksChanged.titleDescription).isNull()
-            assertThat(task2AfterVisibleTasksChanged.title).isNull()
-        }
-
-    @Test
-    fun retrievedThumbnailsCauseEmissionOnTaskDataFlow() =
-        testScope.runTest {
-            // Setup fakes
-            recentsModel.seedTasks(defaultTaskList)
-            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-
-            // Setup TasksRepository
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            val task2DataFlow = systemUnderTest.getTaskDataById(2)
-            val task2BitmapValues = mutableListOf<Bitmap?>()
-            testScope.backgroundScope.launch {
-                task2DataFlow.map { it?.thumbnail?.thumbnail }.toList(task2BitmapValues)
-            }
-
-            // Check for first emission
-            assertThat(task2BitmapValues.single()).isNull()
-
-            systemUnderTest.setVisibleTasks(setOf(2))
-            // Check for second emission
-            assertThat(task2BitmapValues).isEqualTo(listOf(null, bitmap2))
-        }
-
-    @Test
-    fun onTaskThumbnailChanged_setsNewThumbnailDataOnTask() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-            systemUnderTest.setVisibleTasks(setOf(1))
-
-            val expectedThumbnailData = createThumbnailData()
-            val expectedPreviousBitmap = taskThumbnailDataSource.taskIdToBitmap[1]
-            val taskDataFlow = systemUnderTest.getTaskDataById(1)
-
-            val task1ThumbnailValues = mutableListOf<ThumbnailData?>()
-            testScope.backgroundScope.launch {
-                taskDataFlow.map { it?.thumbnail }.toList(task1ThumbnailValues)
-            }
-            taskVisualsChangedDelegate.onTaskThumbnailChanged(1, expectedThumbnailData)
-
-            assertThat(task1ThumbnailValues.first()!!.thumbnail).isEqualTo(expectedPreviousBitmap)
-            assertThat(task1ThumbnailValues.last()).isEqualTo(expectedThumbnailData)
-        }
-
-    @Test
-    fun onHighResLoadingStateChanged_highResReplacesLowResThumbnail() =
-        testScope.runTest {
-            taskThumbnailDataSource.highResEnabled = false
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1))
-
-            val expectedBitmap = mock<Bitmap>()
-            val expectedPreviousBitmap = taskThumbnailDataSource.taskIdToBitmap[1]
-            val taskDataFlow = systemUnderTest.getTaskDataById(1)
-
-            val task1ThumbnailValues = mutableListOf<ThumbnailData?>()
-            testScope.backgroundScope.launch {
-                taskDataFlow.map { it?.thumbnail }.toList(task1ThumbnailValues)
-            }
-
-            taskThumbnailDataSource.taskIdToBitmap[1] = expectedBitmap
-            taskThumbnailDataSource.highResEnabled = true
-            taskVisualsChangedDelegate.onHighResLoadingStateChanged(true)
-
-            val firstThumbnailValue = task1ThumbnailValues.first()!!
-            assertThat(firstThumbnailValue.thumbnail).isEqualTo(expectedPreviousBitmap)
-            assertThat(firstThumbnailValue.reducedResolution).isTrue()
-
-            val lastThumbnailValue = task1ThumbnailValues.last()!!
-            assertThat(lastThumbnailValue.thumbnail).isEqualTo(expectedBitmap)
-            assertThat(lastThumbnailValue.reducedResolution).isFalse()
-        }
-
-    @Test
-    fun onHighResLoadingStateChanged_invisibleTaskIgnored() =
-        testScope.runTest {
-            taskThumbnailDataSource.highResEnabled = false
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1))
-
-            val invisibleTaskId = 2
-            val taskDataFlow = systemUnderTest.getTaskDataById(invisibleTaskId)
-
-            val task2ThumbnailValues = mutableListOf<ThumbnailData?>()
-            testScope.backgroundScope.launch {
-                taskDataFlow.map { it?.thumbnail }.toList(task2ThumbnailValues)
-            }
-
-            taskThumbnailDataSource.highResEnabled = true
-            taskVisualsChangedDelegate.onHighResLoadingStateChanged(true)
-
-            assertThat(task2ThumbnailValues.filterNotNull()).isEmpty()
-            assertThat(taskThumbnailDataSource.getNumberOfGetThumbnailCalls(2)).isEqualTo(0)
-        }
-
-    @Test
-    fun onHighResLoadingStateChanged_lowResDoesNotReplaceHighResThumbnail() =
-        testScope.runTest {
-            taskThumbnailDataSource.highResEnabled = true
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1))
-
-            val expectedBitmap = mock<Bitmap>()
-            val expectedPreviousBitmap = taskThumbnailDataSource.taskIdToBitmap[1]
-            val taskDataFlow = systemUnderTest.getTaskDataById(1)
-
-            val task1ThumbnailValues = mutableListOf<ThumbnailData?>()
-            testScope.backgroundScope.launch {
-                taskDataFlow.map { it?.thumbnail }.toList(task1ThumbnailValues)
-            }
-
-            taskThumbnailDataSource.taskIdToBitmap[1] = expectedBitmap
-            taskThumbnailDataSource.highResEnabled = false
-            taskVisualsChangedDelegate.onHighResLoadingStateChanged(false)
-
-            val firstThumbnailValue = task1ThumbnailValues.first()!!
-            assertThat(firstThumbnailValue.thumbnail).isEqualTo(expectedPreviousBitmap)
-            assertThat(firstThumbnailValue.reducedResolution).isFalse()
-
-            val lastThumbnailValue = task1ThumbnailValues.last()!!
-            assertThat(lastThumbnailValue.thumbnail).isEqualTo(expectedPreviousBitmap)
-            assertThat(lastThumbnailValue.reducedResolution).isFalse()
-        }
-
-    @Test
-    fun onTaskIconChanged_setsNewIconOnTask() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            systemUnderTest.setVisibleTasks(setOf(1))
-
-            val expectedIcon = FakeTaskIconDataSource.mockCopyableDrawable()
-            val expectedPreviousIcon = taskIconDataSource.taskIdToDrawable[1]
-            val taskDataFlow = systemUnderTest.getTaskDataById(1)
-
-            val task1IconValues = mutableListOf<Drawable?>()
-            testScope.backgroundScope.launch {
-                taskDataFlow.map { it?.icon }.toList(task1IconValues)
-            }
-            taskIconDataSource.taskIdToDrawable[1] = expectedIcon
-            taskVisualsChangedDelegate.onTaskIconChanged(1)
-
-            assertThat(task1IconValues.first()).isEqualTo(expectedPreviousIcon)
-            assertThat(task1IconValues.last()).isEqualTo(expectedIcon)
-        }
-
-    @Test
-    fun setVisibleTasks_multipleTimesWithDifferentTasks_reusesThumbnailRequests() =
-        testScope.runTest {
-            recentsModel.seedTasks(defaultTaskList)
-            systemUnderTest.getAllTaskData(forceRefresh = true)
-
-            val taskDataFlow = systemUnderTest.getTaskDataById(1)
-            val task1IconValues = mutableListOf<Drawable?>()
-            testScope.backgroundScope.launch {
-                taskDataFlow.map { it?.icon }.toList(task1IconValues)
-            }
-
-            systemUnderTest.setVisibleTasks(setOf(1))
-            assertThat(taskThumbnailDataSource.getNumberOfGetThumbnailCalls(1)).isEqualTo(1)
-
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
-            assertThat(taskThumbnailDataSource.getNumberOfGetThumbnailCalls(1)).isEqualTo(1)
-        }
+        // Check for second emission
+        assertThat(taskFlowValuesList[1]!!.thumbnail!!.thumbnail).isEqualTo(bitmap2)
+    }
 
     private fun createTaskWithId(taskId: Int) =
         Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000))
-
-    private fun createThumbnailData(): ThumbnailData {
-        val bitmap = mock<Bitmap>()
-        whenever(bitmap.width).thenReturn(THUMBNAIL_WIDTH)
-        whenever(bitmap.height).thenReturn(THUMBNAIL_HEIGHT)
-
-        return ThumbnailData(thumbnail = bitmap)
-    }
-
-    companion object {
-        const val THUMBNAIL_WIDTH = 100
-        const val THUMBNAIL_HEIGHT = 200
-    }
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt

@@ -0,0 +1,178 @@
+/*
+ * 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.task.thumbnail
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Bitmap
+import android.graphics.Color
+import android.graphics.Rect
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
+import com.android.quickstep.task.viewmodel.TaskViewData
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+@RunWith(AndroidJUnit4::class)
+class TaskThumbnailViewModelTest {
+    private val recentsViewData = RecentsViewData()
+    private val taskViewData = TaskViewData()
+    private val tasksRepository = FakeTasksRepository()
+    private val systemUnderTest =
+        TaskThumbnailViewModel(recentsViewData, taskViewData, tasksRepository)
+
+    private val tasks = (0..5).map(::createTaskWithId)
+
+    @Test
+    fun initialStateIsUninitialized() = runTest {
+        assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized)
+    }
+
+    @Test
+    fun bindRunningTask_thenStateIs_LiveTile() = runTest {
+        tasksRepository.seedTasks(tasks)
+        val taskThumbnail = TaskThumbnail(taskId = 1, isRunning = true)
+        systemUnderTest.bind(taskThumbnail)
+
+        assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile)
+    }
+
+    @Test
+    fun setRecentsFullscreenProgress_thenProgressIsPassedThrough() = runTest {
+        recentsViewData.fullscreenProgress.value = 0.5f
+
+        assertThat(systemUnderTest.recentsFullscreenProgress.first()).isEqualTo(0.5f)
+
+        recentsViewData.fullscreenProgress.value = 0.6f
+
+        assertThat(systemUnderTest.recentsFullscreenProgress.first()).isEqualTo(0.6f)
+    }
+
+    @Test
+    fun setAncestorScales_thenScaleIsCalculated() = runTest {
+        recentsViewData.scale.value = 0.5f
+        taskViewData.scale.value = 0.6f
+
+        assertThat(systemUnderTest.inheritedScale.first()).isEqualTo(0.3f)
+    }
+
+    @Test
+    fun bindRunningTaskThenStoppedTaskWithoutThumbnail_thenStateChangesToBackgroundOnly() =
+        runTest {
+            tasksRepository.seedTasks(tasks)
+            val runningTask = TaskThumbnail(taskId = 1, isRunning = true)
+            val stoppedTask = TaskThumbnail(taskId = 2, isRunning = false)
+            systemUnderTest.bind(runningTask)
+            assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile)
+
+            systemUnderTest.bind(stoppedTask)
+            assertThat(systemUnderTest.uiState.first())
+                .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
+        }
+
+    @Test
+    fun bindStoppedTaskWithoutThumbnail_thenStateIs_BackgroundOnly_withAlphaRemoved() = runTest {
+        tasksRepository.seedTasks(tasks)
+        val stoppedTask = TaskThumbnail(taskId = 2, isRunning = false)
+
+        systemUnderTest.bind(stoppedTask)
+        assertThat(systemUnderTest.uiState.first())
+            .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
+    }
+
+    @Test
+    fun bindLockedTaskWithThumbnail_thenStateIs_BackgroundOnly() = runTest {
+        tasksRepository.seedThumbnailData(mapOf(2 to createThumbnailData()))
+        tasks[2].isLocked = true
+        tasksRepository.seedTasks(tasks)
+        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
+
+        systemUnderTest.bind(recentTask)
+        assertThat(systemUnderTest.uiState.first())
+            .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
+    }
+
+    @Test
+    fun bindStoppedTaskWithThumbnail_thenStateIs_Snapshot_withAlphaRemoved() = runTest {
+        val expectedThumbnailData = createThumbnailData()
+        tasksRepository.seedThumbnailData(mapOf(2 to expectedThumbnailData))
+        tasksRepository.seedTasks(tasks)
+        tasksRepository.setVisibleTasks(listOf(2))
+        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
+
+        systemUnderTest.bind(recentTask)
+        assertThat(systemUnderTest.uiState.first())
+            .isEqualTo(
+                Snapshot(
+                    backgroundColor = Color.rgb(2, 2, 2),
+                    bitmap = expectedThumbnailData.thumbnail!!,
+                    drawnRect = Rect(0, 0, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
+                )
+            )
+    }
+
+    @Test
+    fun bindNonVisibleStoppedTask_whenMadeVisible_thenStateIsSnapshot() = runTest {
+        val expectedThumbnailData = createThumbnailData()
+        tasksRepository.seedThumbnailData(mapOf(2 to expectedThumbnailData))
+        tasksRepository.seedTasks(tasks)
+        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
+
+        systemUnderTest.bind(recentTask)
+        assertThat(systemUnderTest.uiState.first())
+            .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
+        tasksRepository.setVisibleTasks(listOf(2))
+        assertThat(systemUnderTest.uiState.first())
+            .isEqualTo(
+                Snapshot(
+                    backgroundColor = Color.rgb(2, 2, 2),
+                    bitmap = expectedThumbnailData.thumbnail!!,
+                    drawnRect = Rect(0, 0, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
+                )
+            )
+    }
+
+    private fun createTaskWithId(taskId: Int) =
+        Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
+            colorBackground = Color.argb(taskId, taskId, taskId, taskId)
+        }
+
+    private fun createThumbnailData(): ThumbnailData {
+        val bitmap = mock<Bitmap>()
+        whenever(bitmap.width).thenReturn(THUMBNAIL_WIDTH)
+        whenever(bitmap.height).thenReturn(THUMBNAIL_HEIGHT)
+
+        return ThumbnailData(thumbnail = bitmap)
+    }
+
+    companion object {
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
+}

quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt

@@ -16,10 +16,7 @@
 
 package com.android.quickstep.util
 
-import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
-import android.content.res.Resources
-import android.view.Display
-import android.view.Display.DEFAULT_DISPLAY
+import android.content.Context
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.apppairs.AppPairIcon
 import com.android.launcher3.logging.StatsLogManager
@@ -27,15 +24,13 @@ import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_BOTTOM_OR_RIGHT
 import com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT
-import com.android.launcher3.views.ActivityContext
 import com.android.quickstep.TopTaskTracker
 import com.android.quickstep.TopTaskTracker.CachedTaskInfo
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.Task.TaskKey
-import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
-import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_33_66
-import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50
-import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_66_33
+import com.android.wm.shell.common.split.SplitScreenConstants.SNAP_TO_30_70
+import com.android.wm.shell.common.split.SplitScreenConstants.SNAP_TO_50_50
+import com.android.wm.shell.common.split.SplitScreenConstants.SNAP_TO_70_30
 import java.util.function.Consumer
 import org.junit.Assert.assertEquals
 import org.junit.Before
@@ -58,34 +53,32 @@ import org.mockito.kotlin.whenever
 
 @RunWith(AndroidJUnit4::class)
 class AppPairsControllerTest {
-    @Mock lateinit var context: ActivityContext
-    @Mock lateinit var resources: Resources
+    @Mock lateinit var context: Context
     @Mock lateinit var splitSelectStateController: SplitSelectStateController
     @Mock lateinit var statsLogManager: StatsLogManager
 
     private lateinit var appPairsController: AppPairsController
 
-    private val left33: Int by lazy {
-        appPairsController.encodeRank(STAGE_POSITION_TOP_OR_LEFT, SNAP_TO_2_33_66)
+    private val left30: Int by lazy {
+        appPairsController.encodeRank(STAGE_POSITION_TOP_OR_LEFT, SNAP_TO_30_70)
     }
     private val left50: Int by lazy {
-        appPairsController.encodeRank(STAGE_POSITION_TOP_OR_LEFT, SNAP_TO_2_50_50)
+        appPairsController.encodeRank(STAGE_POSITION_TOP_OR_LEFT, SNAP_TO_50_50)
     }
-    private val left66: Int by lazy {
-        appPairsController.encodeRank(STAGE_POSITION_TOP_OR_LEFT, SNAP_TO_2_66_33)
+    private val left70: Int by lazy {
+        appPairsController.encodeRank(STAGE_POSITION_TOP_OR_LEFT, SNAP_TO_70_30)
     }
-    private val right33: Int by lazy {
-        appPairsController.encodeRank(STAGE_POSITION_BOTTOM_OR_RIGHT, SNAP_TO_2_33_66)
+    private val right30: Int by lazy {
+        appPairsController.encodeRank(STAGE_POSITION_BOTTOM_OR_RIGHT, SNAP_TO_30_70)
     }
     private val right50: Int by lazy {
-        appPairsController.encodeRank(STAGE_POSITION_BOTTOM_OR_RIGHT, SNAP_TO_2_50_50)
+        appPairsController.encodeRank(STAGE_POSITION_BOTTOM_OR_RIGHT, SNAP_TO_50_50)
     }
-    private val right66: Int by lazy {
-        appPairsController.encodeRank(STAGE_POSITION_BOTTOM_OR_RIGHT, SNAP_TO_2_66_33)
+    private val right70: Int by lazy {
+        appPairsController.encodeRank(STAGE_POSITION_BOTTOM_OR_RIGHT, SNAP_TO_70_30)
     }
 
     @Mock lateinit var mockAppPairIcon: AppPairIcon
-    @Mock lateinit var mockDisplay: Display
     @Mock lateinit var mockTaskbarActivityContext: TaskbarActivityContext
     @Mock lateinit var mockTopTaskTracker: TopTaskTracker
     @Mock lateinit var mockCachedTaskInfo: CachedTaskInfo
@@ -108,42 +101,38 @@ class AppPairsControllerTest {
         // Stub methods on appPairsController so that they return mocks
         spyAppPairsController = spy(appPairsController)
         whenever(mockAppPairIcon.context).thenReturn(mockTaskbarActivityContext)
-        whenever(mockAppPairIcon.display).thenReturn(mockDisplay)
-        whenever(mockDisplay.displayId).thenReturn(DEFAULT_DISPLAY)
         doReturn(mockTopTaskTracker).whenever(spyAppPairsController).topTaskTracker
-        whenever(mockTopTaskTracker.getCachedTopTask(any(), any())).thenReturn(mockCachedTaskInfo)
+        whenever(mockTopTaskTracker.getCachedTopTask(any())).thenReturn(mockCachedTaskInfo)
         whenever(mockTask1.getKey()).thenReturn(mockTaskKey1)
         whenever(mockTask2.getKey()).thenReturn(mockTaskKey2)
         doNothing().whenever(spyAppPairsController).launchAppPair(any(), any())
         doNothing()
             .whenever(spyAppPairsController)
             .launchToSide(anyOrNull(), anyOrNull(), anyOrNull(), anyOrNull())
-        whenever(mockAppPairIcon.context.resources).thenReturn(resources)
-        whenever(DesktopModeStatus.canEnterDesktopMode(mockAppPairIcon.context)).thenReturn(false)
     }
 
     @Test
     fun shouldEncodeRankCorrectly() {
-        assertEquals("left + 33-66 should encode as 0 (0b0)", 0, left33)
+        assertEquals("left + 30-70 should encode as 0 (0b0)", 0, left30)
         assertEquals("left + 50-50 should encode as 1 (0b1)", 1, left50)
-        assertEquals("left + 66-33 should encode as 2 (0b10)", 2, left66)
+        assertEquals("left + 70-30 should encode as 2 (0b10)", 2, left70)
         // See AppPairsController#BITMASK_SIZE and BITMASK_FOR_SNAP_POSITION for context
-        assertEquals("right + 33-66 should encode as 1 followed by 16 0s", 1 shl 16, right33)
+        assertEquals("right + 30-70 should encode as 1 followed by 16 0s", 1 shl 16, right30)
         assertEquals("right + 50-50 should encode as the above value + 1", (1 shl 16) + 1, right50)
-        assertEquals("right + 66-33 should encode as the above value + 2", (1 shl 16) + 2, right66)
+        assertEquals("right + 70-30 should encode as the above value + 2", (1 shl 16) + 2, right70)
     }
 
     @Test
     fun shouldDecodeRankCorrectly() {
         assertEquals(
-            "left + 33-66 should decode to left",
+            "left + 30-70 should decode to left",
             STAGE_POSITION_TOP_OR_LEFT,
-            AppPairsController.convertRankToStagePosition(left33),
+            AppPairsController.convertRankToStagePosition(left30),
         )
         assertEquals(
-            "left + 33-66 should decode to 33-66",
-            SNAP_TO_2_33_66,
-            AppPairsController.convertRankToSnapPosition(left33),
+            "left + 30-70 should decode to 30-70",
+            SNAP_TO_30_70,
+            AppPairsController.convertRankToSnapPosition(left30),
         )
 
         assertEquals(
@@ -153,30 +142,30 @@ class AppPairsControllerTest {
         )
         assertEquals(
             "left + 50-50 should decode to 50-50",
-            SNAP_TO_2_50_50,
+            SNAP_TO_50_50,
             AppPairsController.convertRankToSnapPosition(left50),
         )
 
         assertEquals(
-            "left + 66-33 should decode to left",
+            "left + 70-30 should decode to left",
             STAGE_POSITION_TOP_OR_LEFT,
-            AppPairsController.convertRankToStagePosition(left66),
+            AppPairsController.convertRankToStagePosition(left70),
         )
         assertEquals(
-            "left + 66-33 should decode to 66-33",
-            SNAP_TO_2_66_33,
-            AppPairsController.convertRankToSnapPosition(left66),
+            "left + 70-30 should decode to 70-30",
+            SNAP_TO_70_30,
+            AppPairsController.convertRankToSnapPosition(left70),
         )
 
         assertEquals(
-            "right + 33-66 should decode to right",
+            "right + 30-70 should decode to right",
             STAGE_POSITION_BOTTOM_OR_RIGHT,
-            AppPairsController.convertRankToStagePosition(right33),
+            AppPairsController.convertRankToStagePosition(right30),
         )
         assertEquals(
-            "right + 33-66 should decode to 33-66",
-            SNAP_TO_2_33_66,
-            AppPairsController.convertRankToSnapPosition(right33),
+            "right + 30-70 should decode to 30-70",
+            SNAP_TO_30_70,
+            AppPairsController.convertRankToSnapPosition(right30),
         )
 
         assertEquals(
@@ -186,19 +175,19 @@ class AppPairsControllerTest {
         )
         assertEquals(
             "right + 50-50 should decode to 50-50",
-            SNAP_TO_2_50_50,
+            SNAP_TO_50_50,
             AppPairsController.convertRankToSnapPosition(right50),
         )
 
         assertEquals(
-            "right + 66-33 should decode to right",
+            "right + 70-30 should decode to right",
             STAGE_POSITION_BOTTOM_OR_RIGHT,
-            AppPairsController.convertRankToStagePosition(right66),
+            AppPairsController.convertRankToStagePosition(right70),
         )
         assertEquals(
-            "right + 66-33 should decode to 66-33",
-            SNAP_TO_2_66_33,
-            AppPairsController.convertRankToSnapPosition(right66),
+            "right + 70-30 should decode to 70-30",
+            SNAP_TO_70_30,
+            AppPairsController.convertRankToSnapPosition(right70),
         )
     }
 
@@ -213,7 +202,7 @@ class AppPairsControllerTest {
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -237,7 +226,7 @@ class AppPairsControllerTest {
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -261,7 +250,7 @@ class AppPairsControllerTest {
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -285,7 +274,7 @@ class AppPairsControllerTest {
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -309,7 +298,7 @@ class AppPairsControllerTest {
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -333,7 +322,7 @@ class AppPairsControllerTest {
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -352,16 +341,12 @@ class AppPairsControllerTest {
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 1 already on screen
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
-            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(1))).thenReturn(true)
-        } else {
-            whenever(mockCachedTaskInfo.taskId).thenReturn(1)
-        }
+        whenever(mockCachedTaskInfo.taskId).thenReturn(1)
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -380,16 +365,12 @@ class AppPairsControllerTest {
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 2 already on screen
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
-            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(2))).thenReturn(true)
-        } else {
-            whenever(mockCachedTaskInfo.taskId).thenReturn(2)
-        }
+        whenever(mockCachedTaskInfo.taskId).thenReturn(2)
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -402,84 +383,18 @@ class AppPairsControllerTest {
             .launchToSide(anyOrNull(), anyOrNull(), anyOrNull(), eq(STAGE_POSITION_TOP_OR_LEFT))
     }
 
-    @Test
-    fun handleAppPairLaunchInApp_freeformTask1IsOnScreen_shouldLaunchAppPair() {
-        whenever(DesktopModeStatus.canEnterDesktopMode(mockAppPairIcon.context)).thenReturn(true)
-        /// Test launching apps 1 and 2 from app pair
-        whenever(mockTaskKey1.getId()).thenReturn(1)
-        whenever(mockTaskKey2.getId()).thenReturn(2)
-        // Task 1 is in freeform windowing mode
-        mockTaskKey1.windowingMode = WINDOWING_MODE_FREEFORM
-        // ... and app 1 is already on screen
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
-            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(1))).thenReturn(true)
-        } else {
-            whenever(mockCachedTaskInfo.taskId).thenReturn(1)
-        }
-
-        // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
-        spyAppPairsController.handleAppPairLaunchInApp(
-            mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
-        )
-        verify(splitSelectStateController)
-            .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
-        val callback: Consumer<Array<Task>> = callbackCaptor.value
-        callback.accept(arrayOf(mockTask1, mockTask2))
-
-        // Verify that launchAppPair was called
-        verify(spyAppPairsController, times(1)).launchAppPair(any(), any())
-        verify(spyAppPairsController, never())
-            .launchToSide(anyOrNull(), anyOrNull(), anyOrNull(), anyOrNull())
-    }
-
-    @Test
-    fun handleAppPairLaunchInApp_freeformTask2IsOnScreen_shouldLaunchAppPair() {
-        whenever(DesktopModeStatus.canEnterDesktopMode(mockAppPairIcon.context)).thenReturn(true)
-        /// Test launching apps 1 and 2 from app pair
-        whenever(mockTaskKey1.getId()).thenReturn(1)
-        whenever(mockTaskKey2.getId()).thenReturn(2)
-        // Task 2 is in freeform windowing mode
-        mockTaskKey1.windowingMode = WINDOWING_MODE_FREEFORM
-        // ... and app 2 is already on screen
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
-            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(2))).thenReturn(true)
-        } else {
-            whenever(mockCachedTaskInfo.taskId).thenReturn(2)
-        }
-
-        // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
-        spyAppPairsController.handleAppPairLaunchInApp(
-            mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
-        )
-        verify(splitSelectStateController)
-            .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
-        val callback: Consumer<Array<Task>> = callbackCaptor.value
-        callback.accept(arrayOf(mockTask1, mockTask2))
-
-        // Verify that launchAppPair was called
-        verify(spyAppPairsController, times(1)).launchAppPair(any(), any())
-        verify(spyAppPairsController, never())
-            .launchToSide(anyOrNull(), anyOrNull(), anyOrNull(), anyOrNull())
-    }
-
     @Test
     fun handleAppPairLaunchInApp_shouldLaunchAppPairNormallyWhenUnrelatedSingleAppIsFullscreen() {
         // Test launching apps 1 and 2 from app pair
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 3 already on screen
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
-            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(3))).thenReturn(true)
-        } else {
-            whenever(mockCachedTaskInfo.taskId).thenReturn(3)
-        }
+        whenever(mockCachedTaskInfo.taskId).thenReturn(3)
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2),
+            listOf(mockItemInfo1, mockItemInfo2)
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())

quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt

@@ -27,15 +27,15 @@ import android.view.View
 import android.window.TransitionInfo
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.apppairs.AppPairIcon
-import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.statehandlers.DepthController
 import com.android.launcher3.statemanager.StateManager
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.quickstep.views.GroupedTaskView
 import com.android.quickstep.views.IconView
-import com.android.quickstep.views.TaskContainer
+import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
+import com.android.quickstep.views.TaskView.TaskContainer
 import com.android.systemui.shared.recents.model.Task
 import org.junit.Assert.assertEquals
 import org.junit.Before
@@ -59,7 +59,7 @@ class SplitAnimationControllerTest {
     private val mockSplitSelectStateController: SplitSelectStateController = mock()
     // TaskView
     private val mockTaskView: TaskView = mock()
-    private val mockSnapshotView: View = mock()
+    private val mockThumbnailView: TaskThumbnailViewDeprecated = mock()
     private val mockBitmap: Bitmap = mock()
     private val mockIconView: IconView = mock()
     private val mockTaskViewDrawable: Drawable = mock()
@@ -77,7 +77,6 @@ class SplitAnimationControllerTest {
     private val splitSelectSource: SplitConfigurationOptions.SplitSelectSource = mock()
     private val mockSplitSourceDrawable: Drawable = mock()
     private val mockSplitSourceView: View = mock()
-    private val mockItemInfo: ItemInfo = mock()
 
     private val stateManager: StateManager<*, *> = mock()
     private val depthController: DepthController = mock()
@@ -88,17 +87,14 @@ class SplitAnimationControllerTest {
 
     @Before
     fun setup() {
-        whenever(mockTaskContainer.snapshotView).thenReturn(mockSnapshotView)
-        whenever(mockTaskContainer.thumbnail).thenReturn(mockBitmap)
+        whenever(mockTaskContainer.thumbnailViewDeprecated).thenReturn(mockThumbnailView)
+        whenever(mockThumbnailView.thumbnail).thenReturn(mockBitmap)
         whenever(mockTaskContainer.iconView).thenReturn(mockIconView)
-        whenever(mockTaskContainer.task).thenReturn(mockTask)
         whenever(mockIconView.drawable).thenReturn(mockTaskViewDrawable)
         whenever(mockTaskView.taskContainers).thenReturn(List(1) { mockTaskContainer })
-        whenever(mockTaskView.firstTaskContainer).thenReturn(mockTaskContainer)
 
         whenever(splitSelectSource.drawable).thenReturn(mockSplitSourceDrawable)
         whenever(splitSelectSource.view).thenReturn(mockSplitSourceView)
-        whenever(splitSelectSource.itemInfo).thenReturn(mockItemInfo)
 
         splitAnimationController = SplitAnimationController(mockSplitSelectStateController)
     }
@@ -118,7 +114,7 @@ class SplitAnimationControllerTest {
         assertEquals(
             "Did not fallback to use splitSource icon drawable",
             mockSplitSourceDrawable,
-            splitAnimInitProps.iconDrawable,
+            splitAnimInitProps.iconDrawable
         )
     }
 
@@ -134,7 +130,7 @@ class SplitAnimationControllerTest {
         assertEquals(
             "Did not use taskView icon drawable",
             mockTaskViewDrawable,
-            splitAnimInitProps.iconDrawable,
+            splitAnimInitProps.iconDrawable
         )
     }
 
@@ -153,7 +149,7 @@ class SplitAnimationControllerTest {
         assertEquals(
             "Did not use taskView icon drawable",
             mockTaskViewDrawable,
-            splitAnimInitProps.iconDrawable,
+            splitAnimInitProps.iconDrawable
         )
     }
 
@@ -169,7 +165,7 @@ class SplitAnimationControllerTest {
         assertEquals(
             "Did not use splitSource icon drawable",
             mockSplitSourceDrawable,
-            splitAnimInitProps.iconDrawable,
+            splitAnimInitProps.iconDrawable
         )
     }
 
@@ -184,6 +180,7 @@ class SplitAnimationControllerTest {
 
         whenever(mockTaskContainer.task).thenReturn(mockTask)
         whenever(mockTaskContainer.iconView).thenReturn(mockIconView)
+        whenever(mockTaskContainer.thumbnailViewDeprecated).thenReturn(mockThumbnailView)
         whenever(mockTask.getKey()).thenReturn(mockTaskKey)
         whenever(mockTaskKey.getId()).thenReturn(taskId)
         whenever(mockSplitSelectStateController.initialTaskId).thenReturn(taskId)
@@ -191,13 +188,13 @@ class SplitAnimationControllerTest {
         val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
             splitAnimationController.getFirstAnimInitViews(
                 { mockGroupedTaskView },
-                { splitSelectSource },
+                { splitSelectSource }
             )
 
         assertEquals(
             "Did not use splitSource icon drawable",
             mockSplitSourceDrawable,
-            splitAnimInitProps.iconDrawable,
+            splitAnimInitProps.iconDrawable
         )
     }
 
@@ -215,7 +212,7 @@ class SplitAnimationControllerTest {
                 any(),
                 any(),
                 any(),
-                any(),
+                any()
             )
 
         spySplitAnimationController.playSplitLaunchAnimation(
@@ -230,8 +227,7 @@ class SplitAnimationControllerTest {
             depthController,
             null /* info */,
             null /* t */,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
@@ -244,7 +240,7 @@ class SplitAnimationControllerTest {
                 any(),
                 any(),
                 any(),
-                any(),
+                any()
             )
     }
 
@@ -267,8 +263,7 @@ class SplitAnimationControllerTest {
             depthController,
             transitionInfo,
             transaction,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
@@ -281,7 +276,7 @@ class SplitAnimationControllerTest {
         whenever(mockAppPairIcon.context).thenReturn(mockContextThemeWrapper)
         doNothing()
             .whenever(spySplitAnimationController)
-            .composeIconSplitLaunchAnimator(any(), any(), any(), any(), any())
+            .composeIconSplitLaunchAnimator(any(), any(), any(), any())
         doReturn(-1).whenever(spySplitAnimationController).hasChangesForBothAppPairs(any(), any())
 
         spySplitAnimationController.playSplitLaunchAnimation(
@@ -296,12 +291,11 @@ class SplitAnimationControllerTest {
             depthController,
             transitionInfo,
             transaction,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
-            .composeIconSplitLaunchAnimator(any(), any(), any(), any(), any())
+            .composeIconSplitLaunchAnimator(any(), any(), any(), any())
     }
 
     @Test
@@ -325,8 +319,7 @@ class SplitAnimationControllerTest {
             depthController,
             transitionInfo,
             transaction,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
@@ -353,8 +346,7 @@ class SplitAnimationControllerTest {
             depthController,
             transitionInfo,
             transaction,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
@@ -381,8 +373,7 @@ class SplitAnimationControllerTest {
             depthController,
             transitionInfo,
             transaction,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
@@ -394,7 +385,7 @@ class SplitAnimationControllerTest {
         val spySplitAnimationController = spy(splitAnimationController)
         doNothing()
             .whenever(spySplitAnimationController)
-            .composeFadeInSplitLaunchAnimator(any(), any(), any(), any(), any(), any())
+            .composeFadeInSplitLaunchAnimator(any(), any(), any(), any(), any())
 
         spySplitAnimationController.playSplitLaunchAnimation(
             null /* launchingTaskView */,
@@ -408,11 +399,10 @@ class SplitAnimationControllerTest {
             depthController,
             transitionInfo,
             transaction,
-            {} /* finishCallback */,
-            1f, /* cornerRadius */
+            {} /* finishCallback */
         )
 
         verify(spySplitAnimationController)
-            .composeFadeInSplitLaunchAnimator(any(), any(), any(), any(), any(), any())
+            .composeFadeInSplitLaunchAnimator(any(), any(), any(), any(), any())
     }
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt

@@ -22,6 +22,7 @@ import android.app.PendingIntent
 import android.content.ComponentName
 import android.content.Intent
 import android.graphics.Rect
+import android.os.Handler
 import android.os.UserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.LauncherState
@@ -36,10 +37,9 @@ import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.SystemUiProxy
 import com.android.quickstep.util.SplitSelectStateController.SplitFromDesktopController
-import com.android.quickstep.views.RecentsView
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.systemui.shared.recents.model.Task
-import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50
+import com.android.wm.shell.common.split.SplitScreenConstants.SNAP_TO_50_50
 import java.util.function.Consumer
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
@@ -52,7 +52,6 @@ import org.mockito.Mockito.any
 import org.mockito.Mockito.`when`
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.mock
-import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 
@@ -64,11 +63,11 @@ class SplitSelectStateControllerTest {
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogger = mock()
     private val stateManager: StateManager<LauncherState, StatefulActivity<LauncherState>> = mock()
+    private val handler: Handler = mock()
     private val context: RecentsViewContainer = mock()
     private val recentsModel: RecentsModel = mock()
     private val pendingIntent: PendingIntent = mock()
     private val splitFromDesktopController: SplitFromDesktopController = mock()
-    private val recentsView: RecentsView<*, *> = mock()
 
     private lateinit var splitSelectStateController: SplitSelectStateController
 
@@ -76,7 +75,6 @@ class SplitSelectStateControllerTest {
     private val nonPrimaryUserHandle = UserHandle(ActivityManager.RunningTaskInfo().userId + 10)
 
     private var taskIdCounter = 0
-
     private fun getUniqueId(): Int {
         return ++taskIdCounter
     }
@@ -89,12 +87,13 @@ class SplitSelectStateControllerTest {
         splitSelectStateController =
             SplitSelectStateController(
                 context,
+                handler,
                 stateManager,
                 depthController,
                 statsLogManager,
                 systemUiProxy,
                 recentsModel,
-                null, /*activityBackCallback*/
+                null /*activityBackCallback*/
             )
     }
 
@@ -102,14 +101,14 @@ class SplitSelectStateControllerTest {
     fun activeTasks_noMatchingTasks() {
         val nonMatchingComponent = ComponentKey(ComponentName("no", "match"), primaryUserHandle)
         val groupTask1 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pomegranate", "juice"),
-                ComponentName("pumpkin", "pie"),
+                ComponentName("pumpkin", "pie")
             )
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("hotdog", "juice"),
-                ComponentName("personal", "computer"),
+                ComponentName("personal", "computer")
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask1)
@@ -126,7 +125,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonMatchingComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -143,14 +142,14 @@ class SplitSelectStateControllerTest {
         val matchingComponent =
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
         val groupTask1 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName(matchingPackage, matchingClass),
-                ComponentName("pomegranate", "juice"),
+                ComponentName("pomegranate", "juice")
             )
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pumpkin", "pie"),
-                ComponentName("personal", "computer"),
+                ComponentName("personal", "computer")
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask1)
@@ -163,14 +162,14 @@ class SplitSelectStateControllerTest {
                 assertEquals(
                     "ComponentName package mismatched",
                     it[0].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[0].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
-                assertEquals(it[0], groupTask1.topLeftTask)
+                assertEquals(it[0], groupTask1.task1)
             }
 
         // Capture callback from recentsModel#getTasks()
@@ -179,7 +178,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -196,14 +195,14 @@ class SplitSelectStateControllerTest {
         val nonPrimaryUserComponent =
             ComponentKey(ComponentName(matchingPackage, matchingClass), nonPrimaryUserHandle)
         val groupTask1 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName(matchingPackage, matchingClass),
-                ComponentName("pomegranate", "juice"),
+                ComponentName("pomegranate", "juice")
             )
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pumpkin", "pie"),
-                ComponentName("personal", "computer"),
+                ComponentName("personal", "computer")
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask1)
@@ -220,7 +219,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonPrimaryUserComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -237,16 +236,16 @@ class SplitSelectStateControllerTest {
         val nonPrimaryUserComponent =
             ComponentKey(ComponentName(matchingPackage, matchingClass), nonPrimaryUserHandle)
         val groupTask1 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName(matchingPackage, matchingClass),
                 nonPrimaryUserHandle,
                 ComponentName("pomegranate", "juice"),
-                nonPrimaryUserHandle,
+                nonPrimaryUserHandle
             )
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pumpkin", "pie"),
-                ComponentName("personal", "computer"),
+                ComponentName("personal", "computer")
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask1)
@@ -259,15 +258,15 @@ class SplitSelectStateControllerTest {
                 assertEquals(
                     "ComponentName package mismatched",
                     it[0].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[0].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
                 assertEquals("userId mismatched", it[0].key.userId, nonPrimaryUserHandle.identifier)
-                assertEquals(it[0], groupTask1.topLeftTask)
+                assertEquals(it[0], groupTask1.task1)
             }
 
         // Capture callback from recentsModel#getTasks()
@@ -276,7 +275,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonPrimaryUserComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -293,14 +292,14 @@ class SplitSelectStateControllerTest {
         val matchingComponent =
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
         val groupTask1 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName(matchingPackage, matchingClass),
-                ComponentName("pumpkin", "pie"),
+                ComponentName("pumpkin", "pie")
             )
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pomegranate", "juice"),
-                ComponentName(matchingPackage, matchingClass),
+                ComponentName(matchingPackage, matchingClass)
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask2)
@@ -313,14 +312,14 @@ class SplitSelectStateControllerTest {
                 assertEquals(
                     "ComponentName package mismatched",
                     it[0].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[0].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
-                assertEquals(it[0], groupTask1.topLeftTask)
+                assertEquals(it[0], groupTask1.task1)
             }
 
         // Capture callback from recentsModel#getTasks()
@@ -329,7 +328,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -348,11 +347,11 @@ class SplitSelectStateControllerTest {
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
 
         val groupTask1 =
-            generateSplitTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
+            generateGroupTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pomegranate", "juice"),
-                ComponentName(matchingPackage, matchingClass),
+                ComponentName(matchingPackage, matchingClass)
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask2)
@@ -367,14 +366,14 @@ class SplitSelectStateControllerTest {
                 assertEquals(
                     "ComponentName package mismatched",
                     it[1].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[1].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
-                assertEquals(it[1], groupTask2.bottomRightTask)
+                assertEquals(it[1], groupTask2.task2)
             }
 
         // Capture callback from recentsModel#getTasks()
@@ -383,7 +382,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonMatchingComponent, matchingComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -401,11 +400,11 @@ class SplitSelectStateControllerTest {
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
 
         val groupTask1 =
-            generateSplitTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
+            generateGroupTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pomegranate", "juice"),
-                ComponentName(matchingPackage, matchingClass),
+                ComponentName(matchingPackage, matchingClass)
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask2)
@@ -419,14 +418,14 @@ class SplitSelectStateControllerTest {
                 assertEquals(
                     "ComponentName package mismatched",
                     it[0].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[0].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
-                assertEquals(it[0], groupTask2.bottomRightTask)
+                assertEquals(it[0], groupTask2.task2)
                 assertNull("No tasks should have matched", it[1] /*task*/)
             }
 
@@ -436,7 +435,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent, matchingComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -454,14 +453,14 @@ class SplitSelectStateControllerTest {
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
 
         val groupTask1 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName(matchingPackage, matchingClass),
-                ComponentName("pumpkin", "pie"),
+                ComponentName("pumpkin", "pie")
             )
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("pomegranate", "juice"),
-                ComponentName(matchingPackage, matchingClass),
+                ComponentName(matchingPackage, matchingClass)
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask2)
@@ -475,25 +474,25 @@ class SplitSelectStateControllerTest {
                 assertEquals(
                     "ComponentName package mismatched",
                     it[0].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[0].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
-                assertEquals(it[0], groupTask1.topLeftTask)
+                assertEquals(it[0], groupTask1.task1)
                 assertEquals(
                     "ComponentName package mismatched",
                     it[1].key.baseIntent.component?.packageName,
-                    matchingPackage,
+                    matchingPackage
                 )
                 assertEquals(
                     "ComponentName class mismatched",
                     it[1].key.baseIntent.component?.className,
-                    matchingClass,
+                    matchingClass
                 )
-                assertEquals(it[1], groupTask2.bottomRightTask)
+                assertEquals(it[1], groupTask2.task2)
             }
 
         // Capture callback from recentsModel#getTasks()
@@ -502,7 +501,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent, matchingComponent),
                         false /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -524,16 +523,16 @@ class SplitSelectStateControllerTest {
             ComponentKey(ComponentName(matchingPackage2, matchingClass2), primaryUserHandle)
 
         val groupTask1 =
-            generateSplitTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
+            generateGroupTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
         val groupTask2 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName(matchingPackage2, matchingClass2),
-                ComponentName(matchingPackage, matchingClass),
+                ComponentName(matchingPackage, matchingClass)
             )
         val groupTask3 =
-            generateSplitTask(
+            generateGroupTask(
                 ComponentName("hotdog", "pie"),
-                ComponentName(matchingPackage, matchingClass),
+                ComponentName(matchingPackage, matchingClass)
             )
         val tasks: ArrayList<GroupTask> = ArrayList()
         tasks.add(groupTask3)
@@ -545,7 +544,7 @@ class SplitSelectStateControllerTest {
         val taskConsumer =
             Consumer<Array<Task>> {
                 assertEquals("Expected array length 2", 2, it.size)
-                assertEquals("Found wrong task", it[0], groupTask2.topLeftTask)
+                assertEquals("Found wrong task", it[0], groupTask2.task1)
             }
 
         // Capture callback from recentsModel#getTasks()
@@ -554,7 +553,7 @@ class SplitSelectStateControllerTest {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent2, matchingComponent),
                         true /* findExactPairMatch */,
-                        taskConsumer,
+                        taskConsumer
                     )
                     verify(recentsModel).getTasks(capture())
                 }
@@ -571,7 +570,7 @@ class SplitSelectStateControllerTest {
             -1 /*stagePosition*/,
             ItemInfo(),
             null /*splitEvent*/,
-            10, /*alreadyRunningTask*/
+            10 /*alreadyRunningTask*/
         )
         assertTrue(splitSelectStateController.isSplitSelectActive)
     }
@@ -583,23 +582,21 @@ class SplitSelectStateControllerTest {
             -1 /*stagePosition*/,
             ItemInfo(),
             null /*splitEvent*/,
-            -1, /*alreadyRunningTask*/
+            -1 /*alreadyRunningTask*/
         )
         assertTrue(splitSelectStateController.isSplitSelectActive)
     }
 
     @Test
     fun resetAfterInitial() {
-        whenever(context.getOverviewPanel<RecentsView<*, *>>()).thenReturn(recentsView)
         splitSelectStateController.setInitialTaskSelect(
             Intent() /*intent*/,
             -1 /*stagePosition*/,
             ItemInfo(),
             null /*splitEvent*/,
-            -1,
+            -1
         )
         splitSelectStateController.resetState()
-        verify(recentsView, times(1)).resetDesktopTaskFromSplitSelectState()
         assertFalse(splitSelectStateController.isSplitSelectActive)
     }
 
@@ -625,26 +622,11 @@ class SplitSelectStateControllerTest {
         verify(splitFromDesktopController).onDestroy()
     }
 
-    @Test
-    fun splitSelectStateControllerDestroyed_doNotResetDeskTopTasks() {
-        whenever(context.getOverviewPanel<RecentsView<*, *>>()).thenReturn(recentsView)
-        splitSelectStateController.setInitialTaskSelect(
-            Intent(), /*intent*/
-            -1, /*stagePosition*/
-            ItemInfo(),
-            null, /*splitEvent*/
-            -1,
-        )
-        splitSelectStateController.onDestroy()
-        splitSelectStateController.resetState()
-        verify(recentsView, times(0)).resetDesktopTaskFromSplitSelectState()
-    }
-
-    /** Generates a [SplitTask] with default userId. */
-    private fun generateSplitTask(
+    // Generate GroupTask with default userId.
+    private fun generateGroupTask(
         task1ComponentName: ComponentName,
-        task2ComponentName: ComponentName,
-    ): SplitTask {
+        task2ComponentName: ComponentName
+    ): GroupTask {
         val task1 = Task()
         var taskInfo = ActivityManager.RunningTaskInfo()
         taskInfo.taskId = getUniqueId()
@@ -660,26 +642,20 @@ class SplitSelectStateControllerTest {
         intent.component = task2ComponentName
         taskInfo.baseIntent = intent
         task2.key = Task.TaskKey(taskInfo)
-        return SplitTask(
+        return GroupTask(
             task1,
             task2,
-            SplitConfigurationOptions.SplitBounds(
-                /* leftTopBounds = */ Rect(),
-                /* rightBottomBounds = */ Rect(),
-                /* leftTopTaskId = */ task1.key.id,
-                /* rightBottomTaskId = */ task2.key.id,
-                /* snapPosition = */ SNAP_TO_2_50_50,
-            ),
+            SplitConfigurationOptions.SplitBounds(Rect(), Rect(), -1, -1, SNAP_TO_50_50)
         )
     }
 
-    /** Generates a [SplitTask] with custom user handles. */
-    private fun generateSplitTask(
+    // Generate GroupTask with custom user handles.
+    private fun generateGroupTask(
         task1ComponentName: ComponentName,
         userHandle1: UserHandle,
         task2ComponentName: ComponentName,
-        userHandle2: UserHandle,
-    ): SplitTask {
+        userHandle2: UserHandle
+    ): GroupTask {
         val task1 = Task()
         var taskInfo = ActivityManager.RunningTaskInfo()
         taskInfo.taskId = getUniqueId()
@@ -698,16 +674,10 @@ class SplitSelectStateControllerTest {
         intent.component = task2ComponentName
         taskInfo.baseIntent = intent
         task2.key = Task.TaskKey(taskInfo)
-        return SplitTask(
+        return GroupTask(
             task1,
             task2,
-            SplitConfigurationOptions.SplitBounds(
-                /* leftTopBounds = */ Rect(),
-                /* rightBottomBounds = */ Rect(),
-                /* leftTopTaskId = */ task1.key.id,
-                /* rightBottomTaskId = */ task2.key.id,
-                /* snapPosition = */ SNAP_TO_2_50_50,
-            ),
+            SplitConfigurationOptions.SplitBounds(Rect(), Rect(), -1, -1, SNAP_TO_50_50)
         )
     }
 }

quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.java

@@ -0,0 +1,510 @@
+/*
+ * Copyright (C) 2023 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 static com.android.quickstep.util.TaskGridNavHelper.CLEAR_ALL_PLACEHOLDER_ID;
+import static com.android.quickstep.util.TaskGridNavHelper.INVALID_FOCUSED_TASK_ID;
+
+import static org.junit.Assert.assertEquals;
+
+import com.android.launcher3.util.IntArray;
+
+import org.junit.Test;
+
+public class TaskGridNavHelperTest {
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_pressDown_goesToBottom() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 1;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_DOWN;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 2, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_pressUp_goesToBottom() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 1;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_UP;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 2, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_pressDown_goesToTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 2;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_DOWN;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 1, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_pressUp_goesToTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 2;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_UP;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 1, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_pressLeft_goesLeft() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 1;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 3, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_pressLeft_goesLeft() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 2;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 4, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_secondItem_pressRight_goesRight() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 3;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 1, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_secondItem_pressRight_goesRight() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 4;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 2, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_pressRight_cycleToClearAll() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 1;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", CLEAR_ALL_PLACEHOLDER_ID, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_pressRight_cycleToClearAll() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 2;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", CLEAR_ALL_PLACEHOLDER_ID, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_lastItem_pressLeft_toClearAll() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 5;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", CLEAR_ALL_PLACEHOLDER_ID, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_lastItem_pressLeft_toClearAll() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 6;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", CLEAR_ALL_PLACEHOLDER_ID, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onClearAll_pressLeft_cycleToFirst() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 1, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onClearAll_pressRight_toLastInBottom() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 6, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_withFocused_onFocused_pressLeft_toTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int focusedTaskId = 99;
+        int currentPageTaskViewId = focusedTaskId;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, focusedTaskId);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 1, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_withFocused_onFocused_pressUp_stayOnFocused() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int focusedTaskId = 99;
+        int currentPageTaskViewId = focusedTaskId;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_UP;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, focusedTaskId);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", focusedTaskId, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_withFocused_onFocused_pressDown_stayOnFocused() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int focusedTaskId = 99;
+        int currentPageTaskViewId = focusedTaskId;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_DOWN;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, focusedTaskId);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", focusedTaskId, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_withFocused_onFocused_pressRight_cycleToClearAll() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int focusedTaskId = 99;
+        int currentPageTaskViewId = focusedTaskId;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, focusedTaskId);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", CLEAR_ALL_PLACEHOLDER_ID, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_withFocused_onClearAll_pressLeft_cycleToFocusedTask() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int focusedTaskId = 99;
+        int currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, focusedTaskId);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", focusedTaskId, nextGridPage);
+    }
+
+    @Test
+    public void longerTopRow_noFocused_atEndTopBeyondBottom_pressDown_stayTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5, 7);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 7;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_DOWN;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 7, nextGridPage);
+    }
+
+    @Test
+    public void longerTopRow_noFocused_atEndTopBeyondBottom_pressUp_stayTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5, 7);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 7;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_UP;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 7, nextGridPage);
+    }
+
+    @Test
+    public void longerTopRow_noFocused_atEndBottom_pressLeft_goToTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5, 7);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 6;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_LEFT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 7, nextGridPage);
+    }
+
+    @Test
+    public void longerTopRow_noFocused_atClearAll_pressRight_goToLonger() {
+        IntArray topIds = IntArray.wrap(1, 3, 5, 7);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 7, nextGridPage);
+    }
+
+    @Test
+    public void longerBottomRow_noFocused_atClearAll_pressRight_goToLonger() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6, 7);
+        int currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_RIGHT;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 7, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_pressTab_goesToBottom() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 1;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_TAB;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 2, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_pressTab_goesToNextTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 2;
+        int delta = 1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_TAB;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 3, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onTop_pressTabWithShift_goesToPreviousBottom() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 3;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_TAB;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 2, nextGridPage);
+    }
+
+    @Test
+    public void equalLengthRows_noFocused_onBottom_pressTabWithShift_goesToTop() {
+        IntArray topIds = IntArray.wrap(1, 3, 5);
+        IntArray bottomIds = IntArray.wrap(2, 4, 6);
+        int currentPageTaskViewId = 2;
+        int delta = -1;
+        @TaskGridNavHelper.TASK_NAV_DIRECTION int direction = TaskGridNavHelper.DIRECTION_TAB;
+        boolean cycle = true;
+        TaskGridNavHelper taskGridNavHelper =
+                new TaskGridNavHelper(topIds, bottomIds, INVALID_FOCUSED_TASK_ID);
+
+        int nextGridPage =
+                taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
+
+        assertEquals("Wrong next page returned.", 1, nextGridPage);
+    }
+}

quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java

@@ -35,9 +35,6 @@ import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.dagger.LauncherAppComponent;
-import com.android.launcher3.dagger.LauncherAppSingleton;
-import com.android.launcher3.util.AllModulesMinusWMProxy;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.Info;
 import com.android.launcher3.util.LauncherModelHelper;
@@ -49,9 +46,6 @@ import com.android.launcher3.util.window.WindowManagerProxy;
 import com.android.quickstep.FallbackActivityInterface;
 import com.android.quickstep.util.SurfaceTransaction.MockProperties;
 
-import dagger.BindsInstance;
-import dagger.Component;
-
 import org.hamcrest.Description;
 import org.hamcrest.TypeSafeMatcher;
 import org.junit.Assert;
@@ -165,11 +159,6 @@ public class TaskViewSimulatorTest {
         void verifyNoTransforms() {
             LauncherModelHelper helper = new LauncherModelHelper();
             try {
-                DisplayController mockController = mock(DisplayController.class);
-
-                helper.sandboxContext.initDaggerComponent(
-                        DaggerTaskViewSimulatorTest_TaskViewSimulatorTestComponent.builder()
-                                .bindDisplayController(mockController));
                 int rotation = mDisplaySize.x > mDisplaySize.y
                         ? Surface.ROTATION_90 : Surface.ROTATION_0;
                 CachedDisplayInfo cdi = new CachedDisplayInfo(mDisplaySize, rotation);
@@ -203,14 +192,17 @@ public class TaskViewSimulatorTest {
 
                 DisplayController.Info info = new Info(
                         configurationContext, wmProxy, perDisplayBoundsCache);
+
+                DisplayController mockController = mock(DisplayController.class);
                 when(mockController.getInfo()).thenReturn(info);
+                helper.sandboxContext.putObject(DisplayController.INSTANCE, mockController);
 
                 mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(helper.sandboxContext)
                         .getBestMatch(mAppBounds.width(), mAppBounds.height(), rotation);
                 mDeviceProfile.updateInsets(mLauncherInsets);
 
                 TaskViewSimulator tvs = new TaskViewSimulator(helper.sandboxContext,
-                        FallbackActivityInterface.INSTANCE, false, 0);
+                        FallbackActivityInterface.INSTANCE);
                 tvs.setDp(mDeviceProfile);
 
                 int launcherRotation = info.rotation;
@@ -279,18 +271,4 @@ public class TaskViewSimulatorTest {
             description.appendValue(mExpected);
         }
     }
-
-    @LauncherAppSingleton
-    @Component(modules = {AllModulesMinusWMProxy.class})
-    interface TaskViewSimulatorTestComponent extends LauncherAppComponent {
-
-        @Component.Builder
-        interface Builder extends LauncherAppComponent.Builder {
-
-            @BindsInstance
-            Builder bindDisplayController(DisplayController controller);
-
-            TaskViewSimulatorTestComponent build();
-        }
-    }
 }