Merge "Moving trackpad device tracking to a separate class" into main
diff --git a/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java
index cd08897..fd71151 100644
--- a/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java
@@ -96,7 +96,7 @@
                 itemInfo = apps.data.stream()
                         .filter(info -> user.equals(info.user) && cn.equals(info.componentName))
                         .map(ai -> {
-                            app.getIconCache().getTitleAndIcon(ai, DEFAULT_LOOKUP_FLAG);
+                            app.getIconCache().getTitleAndIcon(ai, mPredictorState.lookupFlag);
                             return ai.makeWorkspaceItem(context);
                         })
                         .findAny()
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index daba0dd..25e1813 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -23,6 +23,7 @@
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.LauncherSettings.Favorites.DESKTOP_ICON_FLAG;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.hybridhotseat.HotseatPredictionModel.convertDataModelToAppTargetBundle;
@@ -60,6 +61,7 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
+import com.android.launcher3.icons.cache.CacheLookupFlag;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.InstanceIdSequence;
@@ -102,14 +104,14 @@
             nonRestorableItem("LAST_SNAPSHOT_TIME_MILLIS", 0L, ENCRYPTED);
 
     @VisibleForTesting
-    final PredictorState mAllAppsState =
-            new PredictorState(CONTAINER_PREDICTION, "all_apps_predictions");
+    final PredictorState mAllAppsState = new PredictorState(
+            CONTAINER_PREDICTION, "all_apps_predictions", DEFAULT_LOOKUP_FLAG);
     @VisibleForTesting
-    final PredictorState mHotseatState =
-            new PredictorState(CONTAINER_HOTSEAT_PREDICTION, "hotseat_predictions");
+    final PredictorState mHotseatState = new PredictorState(
+            CONTAINER_HOTSEAT_PREDICTION, "hotseat_predictions", DESKTOP_ICON_FLAG);
     @VisibleForTesting
-    final PredictorState mWidgetsRecommendationState =
-            new PredictorState(CONTAINER_WIDGETS_PREDICTION, "widgets_prediction");
+    final PredictorState mWidgetsRecommendationState = new PredictorState(
+            CONTAINER_WIDGETS_PREDICTION, "widgets_prediction", DESKTOP_ICON_FLAG);
 
     private final InvariantDeviceProfile mIDP;
     private final AppEventProducer mAppEventProducer;
@@ -153,7 +155,7 @@
 
         WorkspaceItemFactory factory =
                 new WorkspaceItemFactory(mApp, ums, mPmHelper, pinnedShortcuts, numColumns,
-                        state.containerId);
+                        state.containerId, state.lookupFlag);
         FixedContainerItems fci = new FixedContainerItems(state.containerId,
                 state.storage.read(mApp.getContext(), factory, ums.allUsers::get));
         mDataModel.extraItems.put(state.containerId, fci);
@@ -479,13 +481,15 @@
         public final int containerId;
         public final PersistedItemArray<ItemInfo> storage;
         public AppPredictor predictor;
+        public CacheLookupFlag lookupFlag;
 
         private List<AppTarget> mLastTargets;
 
-        PredictorState(int containerId, String storageName) {
+        PredictorState(int containerId, String storageName, CacheLookupFlag lookupFlag) {
             this.containerId = containerId;
             storage = new PersistedItemArray<>(storageName);
             mLastTargets = Collections.emptyList();
+            this.lookupFlag = lookupFlag;
         }
 
         public void destroyPredictor() {
@@ -538,18 +542,20 @@
         private final Map<ShortcutKey, ShortcutInfo> mPinnedShortcuts;
         private final int mMaxCount;
         private final int mContainer;
+        private final CacheLookupFlag mLookupFlag;
 
         private int mReadCount = 0;
 
         protected WorkspaceItemFactory(LauncherAppState appState, UserManagerState ums,
                 PackageManagerHelper pmHelper, Map<ShortcutKey, ShortcutInfo> pinnedShortcuts,
-                int maxCount, int container) {
+                int maxCount, int container, CacheLookupFlag lookupFlag) {
             mAppState = appState;
             mUMS = ums;
             mPmHelper = pmHelper;
             mPinnedShortcuts = pinnedShortcuts;
             mMaxCount = maxCount;
             mContainer = container;
+            mLookupFlag = lookupFlag;
         }
 
         @Nullable
@@ -573,7 +579,7 @@
                             mPmHelper,
                             mUMS.isUserQuiet(user));
                     info.container = mContainer;
-                    mAppState.getIconCache().getTitleAndIcon(info, lai, DEFAULT_LOOKUP_FLAG);
+                    mAppState.getIconCache().getTitleAndIcon(info, lai, mLookupFlag);
                     mReadCount++;
                     return info.makeWorkspaceItem(mAppState.getContext());
                 }
diff --git a/quickstep/src/com/android/quickstep/InputConsumerUtils.kt b/quickstep/src/com/android/quickstep/InputConsumerUtils.kt
index bea3150..558178f 100644
--- a/quickstep/src/com/android/quickstep/InputConsumerUtils.kt
+++ b/quickstep/src/com/android/quickstep/InputConsumerUtils.kt
@@ -56,8 +56,7 @@
 
     @JvmStatic
     fun <S : BaseState<S>, T> newConsumer(
-        baseContext: Context,
-        tisContext: Context,
+        context: Context,
         resetGestureInputConsumer: ResetGestureInputConsumer?,
         overviewComponentObserver: OverviewComponentObserver,
         deviceState: RecentsAnimationDeviceState,
@@ -77,7 +76,7 @@
         val bubbleControllers = tac?.bubbleControllers
         if (bubbleControllers != null && BubbleBarInputConsumer.isEventOnBubbles(tac, event)) {
             val consumer: InputConsumer =
-                BubbleBarInputConsumer(tisContext, bubbleControllers, inputMonitorCompat)
+                BubbleBarInputConsumer(context, bubbleControllers, inputMonitorCompat)
             logInputConsumerSelectionReason(
                 consumer,
                 newCompoundString("event is on bubbles, creating new input consumer"),
@@ -88,7 +87,7 @@
         if (progressProxy != null) {
             val consumer: InputConsumer =
                 ProgressDelegateInputConsumer(
-                    tisContext,
+                    context,
                     taskAnimationManager,
                     gestureState,
                     inputMonitorCompat,
@@ -109,7 +108,7 @@
             if (gestureState.isTrackpadGesture) deviceState.canStartTrackpadGesture()
             else deviceState.canStartSystemGesture()
 
-        if (!get(tisContext).isUserUnlocked) {
+        if (!get(context).isUserUnlocked) {
             val reasonString = newCompoundString("device locked")
             val consumer =
                 if (canStartSystemGesture) {
@@ -117,7 +116,7 @@
                     // launched while device is locked even after exiting direct boot mode (e.g.
                     // camera).
                     createDeviceLockedInputConsumer(
-                        tisContext,
+                        context,
                         resetGestureInputConsumer,
                         deviceState,
                         gestureState,
@@ -148,7 +147,7 @@
                 )
             base =
                 newBaseConsumer<S, T>(
-                    tisContext,
+                    context,
                     resetGestureInputConsumer,
                     overviewComponentObserver,
                     deviceState,
@@ -186,7 +185,7 @@
                 )
                 base =
                     tryCreateAssistantInputConsumer(
-                        tisContext,
+                        context,
                         deviceState,
                         inputMonitorCompat,
                         base,
@@ -213,7 +212,7 @@
                     )
                     base =
                         TaskbarUnstashInputConsumer(
-                            tisContext,
+                            context,
                             base,
                             inputMonitorCompat,
                             tac,
@@ -237,7 +236,7 @@
                 }
             }
 
-            val navHandle = tac?.navHandle ?: SystemUiProxy.INSTANCE[tisContext]
+            val navHandle = tac?.navHandle ?: SystemUiProxy.INSTANCE[context]
             if (
                 canStartSystemGesture &&
                     !previousGestureState.isRecentsAnimationRunning &&
@@ -256,7 +255,7 @@
                 reasonString.append("using NavHandleLongPressInputConsumer")
                 base =
                     NavHandleLongPressInputConsumer(
-                        tisContext,
+                        context,
                         base,
                         inputMonitorCompat,
                         deviceState,
@@ -286,7 +285,7 @@
                             "%ssystem dialog is showing, using SysUiOverlayInputConsumer",
                             SUBSTRING_PREFIX,
                         )
-                base = SysUiOverlayInputConsumer(baseContext, deviceState, inputMonitorCompat)
+                base = SysUiOverlayInputConsumer(context, deviceState, inputMonitorCompat)
             }
 
             if (
@@ -300,7 +299,7 @@
                             "%sTrackpad 3-finger gesture, using TrackpadStatusBarInputConsumer",
                             SUBSTRING_PREFIX,
                         )
-                base = TrackpadStatusBarInputConsumer(baseContext, base, inputMonitorCompat)
+                base = TrackpadStatusBarInputConsumer(context, base, inputMonitorCompat)
             }
 
             if (deviceState.isScreenPinningActive) {
@@ -312,7 +311,7 @@
                         )
                 // Note: we only allow accessibility to wrap this, and it replaces the previous
                 // base input consumer (which should be NO_OP anyway since topTaskLocked == true).
-                base = ScreenPinnedInputConsumer(tisContext, gestureState)
+                base = ScreenPinnedInputConsumer(context, gestureState)
             }
 
             if (deviceState.canTriggerOneHandedAction(event)) {
@@ -323,7 +322,7 @@
                     reasonPrefix,
                     SUBSTRING_PREFIX,
                 )
-                base = OneHandedModeInputConsumer(tisContext, deviceState, base, inputMonitorCompat)
+                base = OneHandedModeInputConsumer(context, deviceState, base, inputMonitorCompat)
             }
 
             if (deviceState.isAccessibilityMenuAvailable) {
@@ -335,7 +334,7 @@
                 )
                 base =
                     AccessibilityInputConsumer(
-                        tisContext,
+                        context,
                         deviceState,
                         gestureState,
                         base,
@@ -362,7 +361,7 @@
                     reasonPrefix,
                     SUBSTRING_PREFIX,
                 )
-                base = OneHandedModeInputConsumer(tisContext, deviceState, base, inputMonitorCompat)
+                base = OneHandedModeInputConsumer(context, deviceState, base, inputMonitorCompat)
             }
         }
         logInputConsumerSelectionReason(base, reasonString)
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 9870ecf..3bfdc21 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -860,7 +860,6 @@
                 mConsumer.onConsumerAboutToBeSwitched();
                 mGestureState = newGestureState;
                 mConsumer = newConsumer(
-                        getBaseContext(),
                         this,
                         mResetGestureInputConsumer,
                         mOverviewComponentObserver,
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index a315775..8a1b211 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -50,14 +50,15 @@
 
     override fun getAllTaskData(forceRefresh: Boolean): Flow<List<Task>> {
         if (forceRefresh) {
-            recentsModel.getTasks { result ->
+            recentsModel.getTasks { newTaskList ->
+                val oldTaskMap = tasks.value
                 val recentTasks =
-                    result
+                    newTaskList
                         .flatMap { groupTask -> groupTask.tasks }
                         .associateBy { it.key.id }
-                        .also { hashMap ->
+                        .also { newTaskMap ->
                             // Clean tasks that are not in the latest group tasks list.
-                            val tasksNoLongerVisible = hashMap.keys.subtract(tasks.value.keys)
+                            val tasksNoLongerVisible = oldTaskMap.keys.subtract(newTaskMap.keys)
                             removeTasks(tasksNoLongerVisible)
 
                             // Use pre-loaded thumbnail data and icon from the previous list.
@@ -66,12 +67,12 @@
                             val cache =
                                 taskRequests.keys
                                     .mapNotNull { key ->
-                                        val task = tasks.value[key] ?: return@mapNotNull null
+                                        val task = oldTaskMap[key] ?: return@mapNotNull null
                                         key to Pair(task.thumbnail, task.icon)
                                     }
                                     .toMap()
 
-                            hashMap.values.forEach { task ->
+                            newTaskMap.values.forEach { task ->
                                 task.thumbnail = task.thumbnail ?: cache[task.key.id]?.first
                                 task.icon = task.icon ?: cache[task.key.id]?.second
                             }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index 20794bf..a952617 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -127,8 +127,9 @@
 
     override fun onDetachedFromWindow() {
         super.onDetachedFromWindow()
+        val scopeToCancel = viewAttachedScope
         recentsCoroutineScope.launch(dispatcherProvider.background) {
-            viewAttachedScope.cancel("TaskThumbnailView detaching from window")
+            scopeToCancel.cancel("TaskThumbnailView detaching from window")
         }
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index 624310b..ee1ec6e 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -37,7 +37,9 @@
 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)
@@ -56,7 +58,7 @@
     private val taskVisualsChangeNotifier = FakeTaskVisualsChangeNotifier()
     private val highResLoadingStateNotifier = FakeHighResLoadingStateNotifier()
     private val taskVisualsChangedDelegate =
-        TaskVisualsChangedDelegateImpl(taskVisualsChangeNotifier, highResLoadingStateNotifier)
+        spy(TaskVisualsChangedDelegateImpl(taskVisualsChangeNotifier, highResLoadingStateNotifier))
 
     private val dispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(dispatcher)
@@ -131,6 +133,29 @@
         }
 
     @Test
+    fun getAllTaskData_clearsPreviouslyLoadedImagesForRemovedTasks() =
+        testScope.runTest {
+            // Setup data
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+
+            // 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)
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 1c4ce74..fa2eb1e 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -20,6 +20,7 @@
 import static android.os.Process.myUserHandle;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.TestUtil.runOnExecutorSync;
@@ -277,7 +278,8 @@
 
     private WidgetsPredictionUpdateTask newWidgetsPredicationTask(List<AppTarget> appTargets) {
         return new WidgetsPredictionUpdateTask(
-                new PredictorState(CONTAINER_WIDGETS_PREDICTION, "test_widgets_prediction"),
+                new PredictorState(CONTAINER_WIDGETS_PREDICTION, "test_widgets_prediction",
+                        DEFAULT_LOOKUP_FLAG),
                 appTargets);
     }
 
diff --git a/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java b/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
index f522a2c..160c578 100644
--- a/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
@@ -478,7 +478,6 @@
         MotionEvent event = MotionEvent.obtain(0, 0, MotionEvent.ACTION_DOWN, 0, 0, 0);
         InputConsumer inputConsumer = newConsumer(
                 mContext,
-                mContext,
                 mResetGestureInputConsumer,
                 mOverviewComponentObserver,
                 mDeviceState,
diff --git a/shared/Android.bp b/shared/Android.bp
new file mode 100644
index 0000000..6d2f407
--- /dev/null
+++ b/shared/Android.bp
@@ -0,0 +1,31 @@
+// 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 {
+    // See: http://go/android-license-faq
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+// Shared between tests and launcher
+android_library {
+    name: "launcher-testing-shared",
+    srcs: [
+        "src/**/*.java",
+        "src/**/*.kt",
+    ],
+    resource_dirs: [],
+    manifest: "AndroidManifest.xml",
+    sdk_version: "current",
+    min_sdk_version: min_launcher3_sdk_version,
+}
diff --git a/tests/multivalentTests/shared/AndroidManifest.xml b/shared/AndroidManifest.xml
similarity index 100%
rename from tests/multivalentTests/shared/AndroidManifest.xml
rename to shared/AndroidManifest.xml
diff --git a/tests/shared/com/android/launcher3/testing/OWNERS b/shared/OWNERS
similarity index 100%
rename from tests/shared/com/android/launcher3/testing/OWNERS
rename to shared/OWNERS
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/ResourceUtils.java b/shared/src/com/android/launcher3/testing/shared/ResourceUtils.java
similarity index 100%
rename from tests/multivalentTests/shared/com/android/launcher3/testing/shared/ResourceUtils.java
rename to shared/src/com/android/launcher3/testing/shared/ResourceUtils.java
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java b/shared/src/com/android/launcher3/testing/shared/TestProtocol.java
similarity index 98%
rename from tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java
rename to shared/src/com/android/launcher3/testing/shared/TestProtocol.java
index 825b52b..4a7471a 100644
--- a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java
+++ b/shared/src/com/android/launcher3/testing/shared/TestProtocol.java
@@ -76,8 +76,9 @@
         }
     }
 
-    public static final String TEST_INFO_REQUEST_FIELD = "request";
     public static final String TEST_INFO_RESPONSE_FIELD = "response";
+    public static final String TEST_INFO_PARAM_INDEX = "index";
+    public static final String TEST_INFO_PARAM_CELL_SPAN = "cell-span";
 
     public static final String REQUEST_HOME_TO_OVERVIEW_SWIPE_HEIGHT =
             "home-to-overview-swipe-height";
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 1d2d161..5dc5d31 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -19,12 +19,15 @@
 import static android.util.Base64.NO_PADDING;
 import static android.util.Base64.NO_WRAP;
 
+import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
+
 import android.database.sqlite.SQLiteDatabase;
 import android.provider.BaseColumns;
 import android.util.Base64;
 
 import androidx.annotation.NonNull;
 
+import com.android.launcher3.icons.cache.CacheLookupFlag;
 import com.android.launcher3.model.data.ItemInfo;
 
 import java.util.LinkedHashMap;
@@ -352,6 +355,11 @@
         public static String getColumns(long profileId) {
             return String.join(", ", getColumnsToTypes(profileId).keySet());
         }
+
+        /**
+         * Lookup flag to be used for items which are visible on the home screen
+         */
+        public static final CacheLookupFlag DESKTOP_ICON_FLAG = DEFAULT_LOOKUP_FLAG;
     }
 
     /**
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 0ebd69f..88a60ea 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -349,10 +349,10 @@
      *
      * @param iconRequestInfos List of IconRequestInfos representing titles and icons to query.
      * @param user UserHandle all the given iconRequestInfos share
-     * @param useLowResIcons whether we should exclude the icon column from the sql results.
+     * @param lookupFlag what flags to use when loading the icon.
      */
     private <T extends ItemInfoWithIcon> Cursor createBulkQueryCursor(
-            List<IconRequestInfo<T>> iconRequestInfos, UserHandle user, boolean useLowResIcons)
+            List<IconRequestInfo<T>> iconRequestInfos, UserHandle user, CacheLookupFlag lookupFlag)
             throws SQLiteException {
         String[] queryParams = Stream.concat(
                 iconRequestInfos.stream()
@@ -365,7 +365,7 @@
                 ",", Collections.nCopies(queryParams.length - 1, "?"));
 
         return iconDb.query(
-                useLowResIcons ? COLUMNS_LOW_RES : COLUMNS_HIGH_RES,
+                toLookupColumns(lookupFlag),
                 COLUMN_COMPONENT
                         + " IN ( " + componentNameQuery + " )"
                         + " AND " + COLUMN_USER + " = ?",
@@ -423,10 +423,11 @@
             List<IconRequestInfo<T>> filteredList,
             Map<ComponentName, List<IconRequestInfo<T>>> duplicateIconRequestsMap) {
         Trace.beginSection("loadIconSubsectionWithDatabase");
+        CacheLookupFlag lookupFlag = DEFAULT_LOOKUP_FLAG.withUseLowRes(sectionKey.second);
         try (Cursor c = createBulkQueryCursor(
                 filteredList,
                 /* user = */ sectionKey.first,
-                /* useLowResIcons = */ sectionKey.second)) {
+                lookupFlag)) {
             // Database title and icon loading
             int componentNameColumnIndex = c.getColumnIndexOrThrow(COLUMN_COMPONENT);
             while (c.moveToNext()) {
@@ -442,7 +443,7 @@
                                 /* user = */ sectionKey.first,
                                 () -> duplicateIconRequests.get(0).launcherActivityInfo,
                                 LauncherActivityCachingLogic.INSTANCE,
-                                DEFAULT_LOOKUP_FLAG.withUseLowRes(sectionKey.second),
+                                lookupFlag,
                                 c);
 
                         for (IconRequestInfo<T> iconRequest : duplicateIconRequests) {
@@ -519,7 +520,7 @@
             @NonNull final PackageItemInfo infoInOut,
             @NonNull CacheLookupFlag lookupFlag) {
         CacheEntry entry = getEntryForPackageLocked(
-                infoInOut.packageName, infoInOut.user, lookupFlag.useLowRes());
+                infoInOut.packageName, infoInOut.user, lookupFlag);
         applyCacheEntry(entry, infoInOut);
         if (infoInOut.widgetCategory == NO_CATEGORY) {
             return;
@@ -593,7 +594,7 @@
 
     @VisibleForTesting
     synchronized boolean isItemInDb(ComponentKey cacheKey) {
-        return getEntryFromDBLocked(cacheKey, new CacheEntry(), false);
+        return getEntryFromDBLocked(cacheKey, new CacheEntry(), DEFAULT_LOOKUP_FLAG);
     }
 
     /**
diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index b936adf..ddbbdc7 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Favorites.DESKTOP_ICON_FLAG;
+
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherActivityInfo;
@@ -194,7 +196,7 @@
                         WorkspaceItemInfo wii = (WorkspaceItemInfo) itemInfo;
                         wii.title = "";
                         wii.bitmap = cache.getDefaultIcon(item.user);
-                        cache.getTitleAndIcon(wii, wii.getMatchingLookupFlag());
+                        cache.getTitleAndIcon(wii, DESKTOP_ICON_FLAG);
                     }
                 }
 
diff --git a/src/com/android/launcher3/model/ItemInstallQueue.java b/src/com/android/launcher3/model/ItemInstallQueue.java
index c4f222f..8acd7e2 100644
--- a/src/com/android/launcher3/model/ItemInstallQueue.java
+++ b/src/com/android/launcher3/model/ItemInstallQueue.java
@@ -18,10 +18,10 @@
 
 import static android.appwidget.AppWidgetManager.EXTRA_APPWIDGET_ID;
 
+import static com.android.launcher3.LauncherSettings.Favorites.DESKTOP_ICON_FLAG;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
-import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
 import static com.android.launcher3.model.data.AppInfo.makeLaunchIntent;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
@@ -310,7 +310,7 @@
                     }
                     LauncherAppState.getInstance(context).getIconCache()
                             .getTitleAndIcon(si, () -> lai,
-                                    DEFAULT_LOOKUP_FLAG.withUsePackageIcon(usePackageIcon));
+                                    DESKTOP_ICON_FLAG.withUsePackageIcon(usePackageIcon));
                     return Pair.create(si, null);
                 }
                 case ITEM_TYPE_DEEP_SHORTCUT: {
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index bc0d776..4e57944 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -23,9 +23,9 @@
 import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 import static com.android.launcher3.LauncherPrefs.SHOULD_SHOW_SMARTSPACE;
+import static com.android.launcher3.LauncherSettings.Favorites.DESKTOP_ICON_FLAG;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.icons.CacheableShortcutInfo.convertShortcutsToCacheableShortcuts;
-import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_HAS_SHORTCUT_PERMISSION;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_PRIVATE_PROFILE_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_CHANGE_PERMISSION;
@@ -603,10 +603,10 @@
                 info.rank = rank;
 
                 if (info instanceof WorkspaceItemInfo wii
-                        && wii.getMatchingLookupFlag().useLowRes()
+                        && wii.getMatchingLookupFlag().isVisuallyLessThan(DESKTOP_ICON_FLAG)
                         && wii.itemType == Favorites.ITEM_TYPE_APPLICATION
                         && verifiers.stream().anyMatch(it -> it.isItemInPreview(info.rank))) {
-                    mIconCache.getTitleAndIcon(wii, DEFAULT_LOOKUP_FLAG);
+                    mIconCache.getTitleAndIcon(wii, DESKTOP_ICON_FLAG);
                 } else if (info instanceof AppPairInfo api) {
                     api.fetchHiResIconsIfNeeded(mIconCache);
                 }
diff --git a/src/com/android/launcher3/model/data/AppPairInfo.kt b/src/com/android/launcher3/model/data/AppPairInfo.kt
index 073d0e0..e620ac9 100644
--- a/src/com/android/launcher3/model/data/AppPairInfo.kt
+++ b/src/com/android/launcher3/model/data/AppPairInfo.kt
@@ -18,9 +18,9 @@
 
 import android.content.Context
 import com.android.launcher3.LauncherSettings
+import com.android.launcher3.LauncherSettings.Favorites.DESKTOP_ICON_FLAG
 import com.android.launcher3.R
 import com.android.launcher3.icons.IconCache
-import com.android.launcher3.icons.cache.CacheLookupFlag.Companion.DEFAULT_LOOKUP_FLAG
 import com.android.launcher3.logger.LauncherAtom
 import com.android.launcher3.views.ActivityContext
 
@@ -82,8 +82,8 @@
     fun fetchHiResIconsIfNeeded(iconCache: IconCache) {
         getAppContents()
             .stream()
-            .filter { it.matchingLookupFlag.useLowRes() }
-            .forEach { member -> iconCache.getTitleAndIcon(member, DEFAULT_LOOKUP_FLAG) }
+            .filter { it.matchingLookupFlag.isVisuallyLessThan(DESKTOP_ICON_FLAG) }
+            .forEach { member -> iconCache.getTitleAndIcon(member, DESKTOP_ICON_FLAG) }
     }
 
     /**
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index aa3f2f2..fb5c8c7 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -55,9 +55,7 @@
 import com.android.launcher3.Workspace;
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.icons.ClockDrawableWrapper;
-import com.android.launcher3.testing.shared.HotseatCellCenterRequest;
 import com.android.launcher3.testing.shared.TestProtocol;
-import com.android.launcher3.testing.shared.WorkspaceCellCenterRequest;
 import com.android.launcher3.util.ActivityLifecycleCallbacksAdapter;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.ResourceBasedOverride;
@@ -265,15 +263,15 @@
                 });
 
             case TestProtocol.REQUEST_WORKSPACE_CELL_CENTER: {
-                final WorkspaceCellCenterRequest request = extra.getParcelable(
-                        TestProtocol.TEST_INFO_REQUEST_FIELD);
+                Rect cellPos = extra.getParcelable(TestProtocol.TEST_INFO_PARAM_CELL_SPAN);
                 return getLauncherUIProperty(Bundle::putParcelable, launcher -> {
                     final Workspace<?> workspace = launcher.getWorkspace();
                     // TODO(b/216387249): allow caller selecting different pages.
                     CellLayout cellLayout = (CellLayout) workspace.getPageAt(
                             workspace.getCurrentPage());
                     final Rect cellRect = getDescendantRectRelativeToDragLayerForCell(launcher,
-                            cellLayout, request.cellX, request.cellY, request.spanX, request.spanY);
+                            cellLayout, cellPos.left, cellPos.top, cellPos.width(),
+                            cellPos.height());
                     return new Point(cellRect.centerX(), cellRect.centerY());
                 });
             }
@@ -292,12 +290,11 @@
             }
 
             case TestProtocol.REQUEST_HOTSEAT_CELL_CENTER: {
-                final HotseatCellCenterRequest request = extra.getParcelable(
-                        TestProtocol.TEST_INFO_REQUEST_FIELD);
+                int cellIndex = extra.getInt(TestProtocol.TEST_INFO_PARAM_INDEX);
                 return getLauncherUIProperty(Bundle::putParcelable, launcher -> {
                     final Hotseat hotseat = launcher.getHotseat();
                     final Rect cellRect = getDescendantRectRelativeToDragLayerForCell(launcher,
-                            hotseat, request.cellInd, /* cellY= */ 0,
+                            hotseat, cellIndex, /* cellY= */ 0,
                             /* spanX= */ 1, /* spanY= */ 1);
                     // TODO(b/234322284): return the real center point.
                     return new Point(cellRect.left + (cellRect.right - cellRect.left) / 3,
diff --git a/tests/Android.bp b/tests/Android.bp
index e4fecc5..f666bba 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -151,19 +151,6 @@
     test_suites: ["general-tests"],
 }
 
-// Shared between tests and launcher
-android_library {
-    name: "launcher-testing-shared",
-    srcs: [
-        "multivalentTests/shared/com/android/launcher3/testing/shared/**/*.java",
-        "multivalentTests/shared/com/android/launcher3/testing/shared/**/*.kt",
-    ],
-    resource_dirs: [],
-    manifest: "multivalentTests/shared/AndroidManifest.xml",
-    sdk_version: "current",
-    min_sdk_version: min_launcher3_sdk_version,
-}
-
 filegroup {
     name: "launcher-testing-helpers-robo",
     srcs: [
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/OWNERS b/tests/multivalentTests/shared/com/android/launcher3/testing/OWNERS
deleted file mode 100644
index a818d5e..0000000
--- a/tests/multivalentTests/shared/com/android/launcher3/testing/OWNERS
+++ /dev/null
@@ -1,5 +0,0 @@
-vadimt@google.com
-sunnygoyal@google.com
-winsonc@google.com
-hyunyoungs@google.com
-mateuszc@google.com
\ No newline at end of file
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/HotseatCellCenterRequest.java b/tests/multivalentTests/shared/com/android/launcher3/testing/shared/HotseatCellCenterRequest.java
deleted file mode 100644
index 7eb035a..0000000
--- a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/HotseatCellCenterRequest.java
+++ /dev/null
@@ -1,99 +0,0 @@
-/*
- * Copyright (C) 2022 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.testing.shared;
-
-import android.os.Parcel;
-import android.os.Parcelable;
-
-/**
- * Request object for querying a hotseat cell region in Rect.
- */
-public class HotseatCellCenterRequest implements TestInformationRequest {
-    public final int cellInd;
-
-    @Override
-    public int describeContents() {
-        return 0;
-    }
-
-    @Override
-    public void writeToParcel(Parcel dest, int flags) {
-        dest.writeInt(cellInd);
-    }
-
-    @Override
-    public String getRequestName() {
-        return TestProtocol.REQUEST_HOTSEAT_CELL_CENTER;
-    }
-
-    public static final Parcelable.Creator<HotseatCellCenterRequest> CREATOR =
-            new Parcelable.Creator<HotseatCellCenterRequest>() {
-
-                @Override
-                public HotseatCellCenterRequest createFromParcel(Parcel source) {
-                    return new HotseatCellCenterRequest(source);
-                }
-
-                @Override
-                public HotseatCellCenterRequest[] newArray(int size) {
-                    return new HotseatCellCenterRequest[size];
-                }
-            };
-
-    private HotseatCellCenterRequest(int cellInd) {
-        this.cellInd = cellInd;
-    }
-
-    private HotseatCellCenterRequest(Parcel in) {
-        this(in.readInt());
-    }
-
-    /**
-     * Create a builder for HotseatCellCenterRequest.
-     *
-     * @return HotseatCellCenterRequest builder.
-     */
-    public static HotseatCellCenterRequest.Builder builder() {
-        return new HotseatCellCenterRequest.Builder();
-    }
-
-    /**
-     * HotseatCellCenterRequest Builder.
-     */
-    public static final class Builder {
-        private int mCellInd;
-
-        private Builder() {
-            mCellInd = 0;
-        }
-
-        /**
-         * Set the index of hotseat cells.
-         */
-        public HotseatCellCenterRequest.Builder setCellInd(int i) {
-            this.mCellInd = i;
-            return this;
-        }
-
-        /**
-         * build the HotseatCellCenterRequest.
-         */
-        public HotseatCellCenterRequest build() {
-            return new HotseatCellCenterRequest(mCellInd);
-        }
-    }
-}
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestInformationRequest.java b/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestInformationRequest.java
deleted file mode 100644
index 38282032..0000000
--- a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestInformationRequest.java
+++ /dev/null
@@ -1,29 +0,0 @@
-/*
- * Copyright (C) 2021 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.testing.shared;
-
-import android.os.Parcelable;
-
-/**
- * A Request sent to TestInformationHandler can implement this interface to carry more information.
- */
-public interface TestInformationRequest extends Parcelable {
-    /**
-     * The name for handler to dispatch request.
-     */
-    String getRequestName();
-}
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java b/tests/multivalentTests/shared/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java
deleted file mode 100644
index e2cd8ea..0000000
--- a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java
+++ /dev/null
@@ -1,138 +0,0 @@
-/*
- * Copyright (C) 2022 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.testing.shared;
-
-import android.os.Parcel;
-import android.os.Parcelable;
-
-/**
- * Request object for querying a workspace cell region in Rect.
- */
-public class WorkspaceCellCenterRequest implements TestInformationRequest {
-    public final int cellX;
-    public final int cellY;
-    public final int spanX;
-    public final int spanY;
-
-    @Override
-    public int describeContents() {
-        return 0;
-    }
-
-    @Override
-    public void writeToParcel(Parcel dest, int flags) {
-        dest.writeInt(cellX);
-        dest.writeInt(cellY);
-        dest.writeInt(spanX);
-        dest.writeInt(spanY);
-    }
-
-    public static final Parcelable.Creator<WorkspaceCellCenterRequest> CREATOR =
-            new Parcelable.Creator<WorkspaceCellCenterRequest>() {
-
-                @Override
-                public WorkspaceCellCenterRequest createFromParcel(Parcel source) {
-                    return new WorkspaceCellCenterRequest(source);
-                }
-
-                @Override
-                public WorkspaceCellCenterRequest[] newArray(int size) {
-                    return new WorkspaceCellCenterRequest[size];
-                }
-            };
-
-    private WorkspaceCellCenterRequest(int cellX, int cellY, int spanX, int spanY) {
-        this.cellX = cellX;
-        this.cellY = cellY;
-        this.spanX = spanX;
-        this.spanY = spanY;
-    }
-
-    private WorkspaceCellCenterRequest(Parcel in) {
-        this(in.readInt(), in.readInt(), in.readInt(), in.readInt());
-    }
-
-    /**
-     * Create a builder for WorkspaceCellRectRequest.
-     *
-     * @return WorkspaceCellRectRequest builder.
-     */
-    public static WorkspaceCellCenterRequest.Builder builder() {
-        return new WorkspaceCellCenterRequest.Builder();
-    }
-
-    @Override
-    public String getRequestName() {
-        return TestProtocol.REQUEST_WORKSPACE_CELL_CENTER;
-    }
-
-    /**
-     * WorkspaceCellRectRequest Builder.
-     */
-    public static final class Builder {
-        private int mCellX;
-        private int mCellY;
-        private int mSpanX;
-        private int mSpanY;
-
-        private Builder() {
-            this.mCellX = 0;
-            this.mCellY = 0;
-            this.mSpanX = 1;
-            this.mSpanY = 1;
-        }
-
-        /**
-         * Set X coordinate of upper left corner expressed as a cell position
-         */
-        public WorkspaceCellCenterRequest.Builder setCellX(int x) {
-            this.mCellX = x;
-            return this;
-        }
-
-        /**
-         * Set Y coordinate of upper left corner expressed as a cell position
-         */
-        public WorkspaceCellCenterRequest.Builder setCellY(int y) {
-            this.mCellY = y;
-            return this;
-        }
-
-        /**
-         * Set span Width in cells
-         */
-        public WorkspaceCellCenterRequest.Builder setSpanX(int x) {
-            this.mSpanX = x;
-            return this;
-        }
-
-        /**
-         * Set span Height in cells
-         */
-        public WorkspaceCellCenterRequest.Builder setSpanY(int y) {
-            this.mSpanY = y;
-            return this;
-        }
-
-        /**
-         * build the WorkspaceCellRectRequest.
-         */
-        public WorkspaceCellCenterRequest build() {
-            return new WorkspaceCellCenterRequest(mCellX, mCellY, mSpanX, mSpanY);
-        }
-    }
-}
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/cache/CacheLookupFlagTest.kt b/tests/multivalentTests/src/com/android/launcher3/icons/cache/CacheLookupFlagTest.kt
index 7b1851d..8218181 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/cache/CacheLookupFlagTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/cache/CacheLookupFlagTest.kt
@@ -78,4 +78,25 @@
         assertTrue(flag.useLowRes())
         assertFalse(flag.usePackageIcon())
     }
+
+    @Test
+    fun `isVisuallyLessThan does not depend on package icon`() {
+        assertFalse(DEFAULT_LOOKUP_FLAG.isVisuallyLessThan(DEFAULT_LOOKUP_FLAG))
+        assertFalse(
+            DEFAULT_LOOKUP_FLAG.withUsePackageIcon().isVisuallyLessThan(DEFAULT_LOOKUP_FLAG)
+        )
+        assertFalse(
+            DEFAULT_LOOKUP_FLAG.isVisuallyLessThan(DEFAULT_LOOKUP_FLAG.withUsePackageIcon())
+        )
+    }
+
+    @Test
+    fun `isVisuallyLessThan depends on low res`() {
+        assertTrue(DEFAULT_LOOKUP_FLAG.withUseLowRes().isVisuallyLessThan(DEFAULT_LOOKUP_FLAG))
+        assertFalse(DEFAULT_LOOKUP_FLAG.isVisuallyLessThan(DEFAULT_LOOKUP_FLAG.withUseLowRes()))
+        assertTrue(
+            DEFAULT_LOOKUP_FLAG.withUseLowRes()
+                .isVisuallyLessThan(DEFAULT_LOOKUP_FLAG.withUsePackageIcon())
+        )
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index fac73d3..bdfe2ab 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -41,6 +41,7 @@
 import android.content.ContentProviderClient;
 import android.content.ContentResolver;
 import android.content.Context;
+import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ProviderInfo;
 import android.content.res.Configuration;
@@ -79,7 +80,6 @@
 import androidx.test.uiautomator.Until;
 
 import com.android.launcher3.testing.shared.ResourceUtils;
-import com.android.launcher3.testing.shared.TestInformationRequest;
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.systemui.shared.system.QuickStepContract;
 
@@ -377,10 +377,8 @@
         }
     }
 
-    Bundle getTestInfo(TestInformationRequest request) {
-        Bundle extra = new Bundle();
-        extra.putParcelable(TestProtocol.TEST_INFO_REQUEST_FIELD, request);
-        return getTestInfo(request.getRequestName(), null, extra);
+    Bundle getTestInfo(Intent request) {
+        return getTestInfo(request.getAction(), null, request.getExtras());
     }
 
     Insets getTargetInsets() {
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index a29362f..d615879 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -29,6 +29,7 @@
 import static junit.framework.TestCase.assertNotNull;
 import static junit.framework.TestCase.assertTrue;
 
+import android.content.Intent;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.SystemClock;
@@ -44,9 +45,7 @@
 import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
 
-import com.android.launcher3.testing.shared.HotseatCellCenterRequest;
 import com.android.launcher3.testing.shared.TestProtocol;
-import com.android.launcher3.testing.shared.WorkspaceCellCenterRequest;
 
 import java.util.List;
 import java.util.Map;
@@ -499,20 +498,23 @@
     }
 
     static Point getCellCenter(LauncherInstrumentation launcher, int cellX, int cellY) {
-        return launcher.getTestInfo(WorkspaceCellCenterRequest.builder().setCellX(cellX).setCellY(
-                cellY).build()).getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+        return getCellCenter(launcher, cellX, cellY, 1, 1);
     }
 
     static Point getCellCenter(LauncherInstrumentation launcher, int cellX, int cellY, int spanX,
             int spanY) {
-        return launcher.getTestInfo(WorkspaceCellCenterRequest.builder().setCellX(cellX)
-                .setCellY(cellY).setSpanX(spanX).setSpanY(spanY).build())
+        return launcher.getTestInfo(
+                new Intent(TestProtocol.REQUEST_WORKSPACE_CELL_CENTER)
+                        .putExtra(TestProtocol.TEST_INFO_PARAM_CELL_SPAN,
+                                new Rect(cellX, cellY, cellX + spanX, cellY + spanY)))
                 .getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
-    static Point getHotseatCellCenter(LauncherInstrumentation launcher, int cellInd) {
-        return launcher.getTestInfo(HotseatCellCenterRequest.builder()
-                .setCellInd(cellInd).build()).getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    static Point getHotseatCellCenter(LauncherInstrumentation launcher, int cellIndex) {
+        return launcher.getTestInfo(
+                new Intent(TestProtocol.REQUEST_HOTSEAT_CELL_CENTER)
+                        .putExtra(TestProtocol.TEST_INFO_PARAM_INDEX, cellIndex))
+                .getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
     /** Returns the number of rows and columns in the workspace */