Merge "Loading generated preview on-demand instead of keeping everything in memory" into main
diff --git a/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt b/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt
index ac1ffa6..8b064d3 100644
--- a/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt
+++ b/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt
@@ -31,6 +31,8 @@
 import com.android.quickstep.SystemUiProxy
 import com.android.quickstep.TaskViewUtils
 import com.android.quickstep.views.DesktopTaskView
+import com.android.quickstep.views.TaskContainer
+import com.android.quickstep.views.TaskView
 import com.android.window.flags.Flags
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 import java.util.function.Consumer
@@ -62,8 +64,12 @@
     }
 
     /** Launch desktop tasks from recents view */
-    fun moveToDesktop(taskId: Int, transitionSource: DesktopModeTransitionSource) {
-        systemUiProxy.moveToDesktop(taskId, transitionSource)
+    fun moveToDesktop(taskContainer: TaskContainer, transitionSource: DesktopModeTransitionSource) {
+        systemUiProxy.moveToDesktop(
+            taskContainer.task.key.id,
+            transitionSource,
+            /* transition = */ null,
+        )
     }
 
     /** Move task to external display from recents view */
@@ -72,7 +78,7 @@
     }
 
     private class RemoteDesktopLaunchTransitionRunner(
-        private val desktopTaskView: DesktopTaskView,
+        private val taskView: TaskView,
         private val animated: Boolean,
         private val stateManager: StateManager<*, *>,
         private val depthController: DepthController?,
@@ -99,7 +105,7 @@
             MAIN_EXECUTOR.execute {
                 val animator =
                     TaskViewUtils.composeRecentsDesktopLaunchAnimator(
-                        desktopTaskView,
+                        taskView,
                         stateManager,
                         depthController,
                         info,
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
index df00696..3e3f569 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
@@ -224,12 +224,13 @@
 
     override fun isBubbleBarVisible(): Boolean = bubbleBarViewController.hasBubbles() && !isStashed
 
-    override fun onNewBubbleAnimationInterrupted(isStashed: Boolean, bubbleBarTranslationY: Float) =
+    override fun onNewBubbleAnimationInterrupted(isStashed: Boolean, bubbleBarTranslationY: Float) {
         if (isStashed) {
             stashBubbleBarImmediate()
         } else {
             showBubbleBarImmediate(bubbleBarTranslationY)
         }
+    }
 
     /** Check if [ev] belongs to the stash handle or the bubble bar views. */
     override fun isEventOverBubbleBarViews(ev: MotionEvent): Boolean {
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 2991e64..6a25ecb 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -1478,10 +1478,11 @@
     }
 
     /** Call shell to move a task with given `taskId` to desktop  */
-    public void moveToDesktop(int taskId, DesktopModeTransitionSource transitionSource) {
+    public void moveToDesktop(int taskId, DesktopModeTransitionSource transitionSource,
+            @Nullable RemoteTransition transition) {
         if (mDesktopMode != null) {
             try {
-                mDesktopMode.moveToDesktop(taskId, transitionSource);
+                mDesktopMode.moveToDesktop(taskId, transitionSource, transition);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call moveToDesktop", e);
             }
diff --git a/quickstep/src/com/android/quickstep/TaskViewUtils.java b/quickstep/src/com/android/quickstep/TaskViewUtils.java
index 084cede..ecb17f6 100644
--- a/quickstep/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskViewUtils.java
@@ -562,7 +562,7 @@
      * Start recents to desktop animation
      */
     public static AnimatorSet composeRecentsDesktopLaunchAnimator(
-            @NonNull DesktopTaskView launchingTaskView,
+            @NonNull TaskView launchingTaskView,
             @NonNull StateManager stateManager, @Nullable DepthController depthController,
             @NonNull TransitionInfo transitionInfo,
             SurfaceControl.Transaction t, @NonNull Runnable finishCallback) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index da5411b..ae906de 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -6787,8 +6787,8 @@
         if (mDesktopRecentsTransitionController == null) {
             return;
         }
-        mDesktopRecentsTransitionController.moveToDesktop(taskContainer.getTask().key.id,
-                transitionSource);
+
+        mDesktopRecentsTransitionController.moveToDesktop(taskContainer, transitionSource);
         successCallback.run();
     }
 
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index c0bd956..21aabfa 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -300,10 +300,6 @@
     <!-- the distance an icon must be dragged before button drop targets accept it -->
     <dimen name="drag_distanceThreshold">30dp</dimen>
 
-    <!-- Elevation for the drag view. It should be larger than elevation of all other drag sources
-         and drop targets like all-apps and folders -->
-    <dimen name="drag_elevation">30dp</dimen>
-
     <dimen name="drag_flingToDeleteMinVelocity">-1500dp</dimen>
 
     <dimen name="spring_loaded_panel_border">2dp</dimen>
diff --git a/src/com/android/launcher3/CheckLongPressHelper.java b/src/com/android/launcher3/CheckLongPressHelper.java
index 3e4e96b..e9cd16c 100644
--- a/src/com/android/launcher3/CheckLongPressHelper.java
+++ b/src/com/android/launcher3/CheckLongPressHelper.java
@@ -151,6 +151,8 @@
                 handled = mView.performLongClick();
             }
             if (handled) {
+                // Cancel any default long-press action on the view
+                mView.cancelLongPress();
                 mView.setPressed(false);
                 mHasPerformedLongPress = true;
             }
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 5becdde..dddc43f 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static com.android.launcher3.LauncherPrefs.DB_FILE;
 import static com.android.launcher3.LauncherPrefs.FIXED_LANDSCAPE_MODE;
 import static com.android.launcher3.LauncherPrefs.GRID_NAME;
 import static com.android.launcher3.Utilities.dpiFromPx;
@@ -53,6 +54,7 @@
 
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.DotRenderer;
+import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.DeviceGridState;
 import com.android.launcher3.provider.RestoreDbTask;
 import com.android.launcher3.testing.shared.ResourceUtils;
@@ -86,6 +88,8 @@
     public static final MainThreadInitializedObject<InvariantDeviceProfile> INSTANCE =
             new MainThreadInitializedObject<>(InvariantDeviceProfile::new);
 
+    public static final String GRID_NAME_PREFS_KEY = "idp_grid_name";
+
     @Retention(RetentionPolicy.SOURCE)
     @IntDef({TYPE_PHONE, TYPE_MULTI_DISPLAY, TYPE_TABLET})
     public @interface DeviceType {
@@ -241,8 +245,11 @@
     @TargetApi(23)
     private InvariantDeviceProfile(Context context) {
         String gridName = getCurrentGridName(context);
+        FileLog.d(TAG, "New InvariantDeviceProfile, before initGrid(): "
+                + "gridName:" + gridName
+                + ", LauncherPrefs GRID_NAME:" + LauncherPrefs.get(context).get(GRID_NAME)
+                + ", LauncherPrefs DB_FILE:" + LauncherPrefs.get(context).get(DB_FILE));
         initGrid(context, gridName);
-
         DisplayController.INSTANCE.get(context).setPriorityListener(
                 (displayContext, info, flags) -> {
                     if ((flags & (CHANGE_DENSITY | CHANGE_SUPPORTED_BOUNDS
@@ -371,6 +378,10 @@
         }
 
         initGrid(context, displayInfo, displayOption);
+        FileLog.d(TAG, "initGrid: "
+                + "gridName:" + gridName
+                + ", LauncherPrefs GRID_NAME:" + LauncherPrefs.get(context).get(GRID_NAME)
+                + ", LauncherPrefs DB_FILE:" + LauncherPrefs.get(context).get(DB_FILE));
         return displayOption.grid.name;
     }
 
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 01d0a74..a53238d 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -20,8 +20,12 @@
 import static android.content.Context.RECEIVER_EXPORTED;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.InvariantDeviceProfile.GRID_NAME_PREFS_KEY;
+import static com.android.launcher3.LauncherPrefs.DB_FILE;
+import static com.android.launcher3.LauncherPrefs.GRID_NAME;
 import static com.android.launcher3.LauncherPrefs.ICON_STATE;
 import static com.android.launcher3.LauncherPrefs.THEMED_ICONS;
+import static com.android.launcher3.model.DeviceGridState.KEY_DB_FILE;
 import static com.android.launcher3.model.LoaderTask.SMARTSPACE_ON_HOME_SCREEN;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
@@ -47,6 +51,7 @@
 import com.android.launcher3.icons.IconProvider;
 import com.android.launcher3.icons.LauncherIconProvider;
 import com.android.launcher3.icons.LauncherIcons;
+import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.ModelLauncherCallbacks;
 import com.android.launcher3.model.WidgetsFilterDataProvider;
 import com.android.launcher3.notification.NotificationListener;
@@ -70,6 +75,7 @@
 
 public class LauncherAppState implements SafeCloseable {
 
+    public static final String TAG = "LauncherAppState";
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
 
     // We do not need any synchronization for this variable as its only written on UI thread.
@@ -293,6 +299,12 @@
             if (Themes.KEY_THEMED_ICONS.equals(key)) {
                 mIconProvider.setIconThemeSupported(Themes.isThemedIconEnabled(mContext));
                 verifyIconChanged();
+            } else if (GRID_NAME_PREFS_KEY.equals(key)) {
+                FileLog.d(TAG, "onPrefChanged GRID_NAME changed: "
+                        + LauncherPrefs.get(mContext).get(GRID_NAME));
+            } else if (KEY_DB_FILE.equals(key)) {
+                FileLog.d(TAG, "onPrefChanged DB_FILE changed: "
+                        + LauncherPrefs.get(mContext).get(DB_FILE));
             }
         }
     }
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 712c56c..5b9c2fa 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -20,6 +20,7 @@
 import android.content.SharedPreferences
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.BuildConfig.WIDGET_ON_FIRST_SCREEN
+import com.android.launcher3.InvariantDeviceProfile.GRID_NAME_PREFS_KEY
 import com.android.launcher3.LauncherFiles.DEVICE_PREFERENCES_KEY
 import com.android.launcher3.LauncherFiles.SHARED_PREFERENCES_KEY
 import com.android.launcher3.model.DeviceGridState
@@ -138,7 +139,7 @@
         @JvmField
         val GRID_NAME =
             ConstantItem(
-                "idp_grid_name",
+                GRID_NAME_PREFS_KEY,
                 isBackedUp = true,
                 defaultValue = null,
                 encryptionType = EncryptionType.ENCRYPTED,
diff --git a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
index 8505a6d..b0001af 100644
--- a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
@@ -761,6 +761,22 @@
         }
     }
 
+    @Override
+    public void addChildrenForAccessibility(ArrayList<View> arrayList) {
+        super.addChildrenForAccessibility(arrayList);
+        if (!Flags.floatingSearchBar()) {
+            // Searchbox container is visually at the top of the all apps UI but it's present in
+            // end of the children list.
+            // We need to move the searchbox to the top in a11y tree for a11y services to read the
+            // all apps screen in same as visual order.
+            arrayList.stream().filter(v -> v.getId() == R.id.search_container_all_apps)
+                    .findFirst().ifPresent(v -> {
+                        arrayList.remove(v);
+                        arrayList.add(0, v);
+                    });
+        }
+    }
+
     protected void updateHeaderScroll(int scrolledOffset) {
         float prog1 = Utilities.boundToRange((float) scrolledOffset / mHeaderThreshold, 0f, 1f);
         int headerColor = getHeaderColor(prog1);
diff --git a/src/com/android/launcher3/allapps/WorkProfileManager.java b/src/com/android/launcher3/allapps/WorkProfileManager.java
index 6ebab5a..6d7d193 100644
--- a/src/com/android/launcher3/allapps/WorkProfileManager.java
+++ b/src/com/android/launcher3/allapps/WorkProfileManager.java
@@ -148,7 +148,7 @@
         if (getAH() != null) {
             getAH().applyPadding();
         }
-        mWorkUtilityView.setOnClickListener(this::onWorkFabClicked);
+        mWorkUtilityView.getWorkFAB().setOnClickListener(this::onWorkFabClicked);
         return true;
     }
     /**
diff --git a/src/com/android/launcher3/allapps/WorkUtilityView.java b/src/com/android/launcher3/allapps/WorkUtilityView.java
index 5949b78..b263639 100644
--- a/src/com/android/launcher3/allapps/WorkUtilityView.java
+++ b/src/com/android/launcher3/allapps/WorkUtilityView.java
@@ -25,6 +25,7 @@
 import android.graphics.Rect;
 import android.text.TextUtils;
 import android.util.AttributeSet;
+import android.view.View;
 import android.view.ViewGroup;
 import android.view.WindowInsets;
 import android.widget.ImageButton;
@@ -83,6 +84,7 @@
     // Threshold when user scrolls up/down to determine when should button extend/collapse
     private final int mScrollThreshold;
     private ValueAnimator mPauseFABAnim;
+    private View mWorkFAB;
     private TextView mPauseText;
     private ImageView mWorkIcon;
     private ImageButton mSchedulerButton;
@@ -117,6 +119,7 @@
 
         mPauseText = findViewById(R.id.pause_text);
         mWorkIcon = findViewById(R.id.work_icon);
+        mWorkFAB = findViewById(R.id.work_mode_toggle);
         mSchedulerButton = findViewById(R.id.work_scheduler);
         setSelected(true);
         KeyboardInsetAnimationCallback keyboardInsetAnimationCallback =
@@ -388,6 +391,10 @@
         return mScrollThreshold;
     }
 
+    public View getWorkFAB() {
+        return mWorkFAB;
+    }
+
     public void updateStringFromCache(){
         StringCache cache = mActivityContext.getStringCache();
         if (cache != null) {
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index bcee442..67fe889 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -224,7 +224,6 @@
         measure(makeMeasureSpec(width, EXACTLY), makeMeasureSpec(height, EXACTLY));
 
         mBlurSizeOutline = getResources().getDimensionPixelSize(R.dimen.blur_size_medium_outline);
-        setElevation(getResources().getDimension(R.dimen.drag_elevation));
         setWillNotDraw(false);
     }
 
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 8db981f..f56888b 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -124,11 +124,13 @@
         LauncherPrefs.get(context).removeSync(RESTORE_DEVICE);
 
         DeviceGridState deviceGridState = new DeviceGridState(context);
+        FileLog.d(TAG, "restoreIfNeeded: deviceGridState from context: " + deviceGridState);
         String oldPhoneFileName = deviceGridState.getDbFile();
         List<String> previousDbs = existingDbs(context);
         removeOldDBs(context, oldPhoneFileName);
         // The idp before this contains data about the old phone, after this it becomes the idp
         // of the current phone.
+        FileLog.d(TAG, "Resetting IDP to default for restore dest device");
         idp.reset(context);
         trySettingPreviousGridAsCurrent(context, idp, oldPhoneFileName, previousDbs);
     }
@@ -144,17 +146,24 @@
                 context, oldPhoneDbFileName);
         // The grid option could be null if current phone doesn't support the previous db.
         if (oldPhoneGridOption != null) {
+            FileLog.d(TAG, "trySettingPreviousGridAsCurrent:"
+                    + ", oldPhoneDbFileName: " + oldPhoneDbFileName
+                    + ", oldPhoneGridOption: " + oldPhoneGridOption
+                    + ", previousDbs: " + previousDbs);
 
             /* If the user only used the default db on the previous phone and the new default db is
              * bigger than or equal to the previous one, then keep the new default db */
             if (previousDbs.size() == 1 && oldPhoneGridOption.numColumns <= idp.numColumns
                     && oldPhoneGridOption.numRows <= idp.numRows) {
                 /* Keep the user in default grid */
+                FileLog.d(TAG, "Keeping default db from restore as current grid");
                 return;
             }
             /*
              * Here we are setting the previous db as the current one.
              */
+            FileLog.d(TAG, "Setting grid from old device as current grid: "
+                + "oldPhoneGridOption:" + oldPhoneGridOption.name);
             idp.setCurrentGrid(context, oldPhoneGridOption.name);
         }
     }
diff --git a/tests/multivalentTests/src/com/android/launcher3/folder/FolderNameProviderTest.java b/tests/multivalentTests/src/com/android/launcher3/folder/FolderNameProviderTest.java
index 9c15309..b6fa81c 100644
--- a/tests/multivalentTests/src/com/android/launcher3/folder/FolderNameProviderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/folder/FolderNameProviderTest.java
@@ -23,6 +23,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.os.Process;
 import android.os.UserHandle;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -49,16 +50,17 @@
     @Before
     public void setUp() {
         mContext = new ActivityContextWrapper(getApplicationContext());
+        int workUserId = Process.myUserHandle().hashCode() + 1;
         mItem1 = new WorkspaceItemInfo(new AppInfo(
                 new ComponentName("a.b.c", "a.b.c/a.b.c.d"),
                 "title1",
-                UserHandle.of(10),
+                UserHandle.of(workUserId),
                 new Intent().setComponent(new ComponentName("a.b.c", "a.b.c/a.b.c.d"))
         ));
         mItem2 = new WorkspaceItemInfo(new AppInfo(
                 new ComponentName("a.b.c", "a.b.c/a.b.c.d"),
                 "title2",
-                UserHandle.of(10),
+                UserHandle.of(workUserId),
                 new Intent().setComponent(new ComponentName("a.b.c", "a.b.c/a.b.c.d"))
         ));
     }
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt
index 9cc380e..db37247 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt
@@ -20,7 +20,7 @@
 import android.content.ComponentName
 import android.content.Intent
 import android.content.pm.PackageInstaller.SessionInfo
-import android.os.UserHandle
+import android.os.Process.myUserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP
@@ -77,16 +77,18 @@
             SessionInfo().apply {
                 installerPackageName = expectedInstallerPackage
                 appPackageName = expectedAppPackage
+                userId = myUserHandle().identifier
             }
         val sessionInfoUnexpected =
             SessionInfo().apply {
                 installerPackageName = expectedInstallerPackage
                 appPackageName = unexpectedAppPackage
+                userId = myUserHandle().identifier
             }
         val sessionInfoMap: HashMap<PackageUserKey, SessionInfo> =
             hashMapOf(
-                PackageUserKey(unexpectedAppPackage, UserHandle(0)) to sessionInfoExpected,
-                PackageUserKey(expectedAppPackage, UserHandle(0)) to sessionInfoUnexpected,
+                PackageUserKey(unexpectedAppPackage, myUserHandle()) to sessionInfoExpected,
+                PackageUserKey(expectedAppPackage, myUserHandle()) to sessionInfoUnexpected,
             )
 
         // When
@@ -200,16 +202,18 @@
             SessionInfo().apply {
                 installerPackageName = expectedInstallerPackage
                 appPackageName = expectedAppPackage
+                userId = myUserHandle().identifier
             }
         val sessionInfoUnexpected =
             SessionInfo().apply {
                 installerPackageName = expectedInstallerPackage
                 appPackageName = unexpectedAppPackage
+                userId = myUserHandle().identifier
             }
         val sessionInfoMap: HashMap<PackageUserKey, SessionInfo> =
             hashMapOf(
-                PackageUserKey(unexpectedAppPackage, UserHandle(0)) to sessionInfoExpected,
-                PackageUserKey(expectedAppPackage, UserHandle(0)) to sessionInfoUnexpected,
+                PackageUserKey(unexpectedAppPackage, myUserHandle()) to sessionInfoExpected,
+                PackageUserKey(expectedAppPackage, myUserHandle()) to sessionInfoUnexpected,
             )
         val expectedItemInfo = WorkspaceItemInfo().apply { intent = expectedIntent }
         val expectedFolderInfo = FolderInfo().apply { add(expectedItemInfo) }
diff --git a/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java b/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java
index 34b292c..f494f38 100644
--- a/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java
+++ b/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.compat;
 
+import static android.os.Process.myUserHandle;
+
 import static com.android.launcher3.Flags.FLAG_ENABLE_SUPPORT_FOR_ARCHIVING;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -136,7 +138,8 @@
     public void testPromiseIcon_addedArchivedApp() throws Throwable {
         installDummyAppAndWaitForUIUpdate();
         assertThat(executeShellCommand(
-                String.format("pm archive %s", DUMMY_PACKAGE)))
+                String.format("pm archive --user %d %s",
+                        myUserHandle().getIdentifier(), DUMMY_PACKAGE)))
                 .isEqualTo("Success\n");
 
         // Create and add test session
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index 882061f..34cbdee 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -2,6 +2,7 @@
 
 import android.appwidget.AppWidgetManager
 import android.content.Intent
+import android.os.Process
 import android.os.UserHandle
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
@@ -136,7 +137,7 @@
     @Test
     fun loadsDataProperly() =
         with(BgDataModel()) {
-            val MAIN_HANDLE = UserHandle.of(0)
+            val MAIN_HANDLE = Process.myUserHandle()
             val mockUserHandles = arrayListOf<UserHandle>(MAIN_HANDLE)
             `when`(userCache.userProfiles).thenReturn(mockUserHandles)
             `when`(userCache.getUserInfo(MAIN_HANDLE)).thenReturn(UserIconInfo(MAIN_HANDLE, 1))
@@ -186,7 +187,7 @@
     fun setsQuietModeFlagCorrectlyForWorkProfile() =
         with(BgDataModel()) {
             setFlagsRule.enableFlags(Flags.FLAG_ENABLE_PRIVATE_SPACE)
-            val MAIN_HANDLE = UserHandle.of(0)
+            val MAIN_HANDLE = Process.myUserHandle()
             val mockUserHandles = arrayListOf<UserHandle>(MAIN_HANDLE)
             `when`(userCache.userProfiles).thenReturn(mockUserHandles)
             `when`(userManagerState?.isUserQuiet(MAIN_HANDLE)).thenReturn(true)
@@ -215,7 +216,7 @@
     fun setsQuietModeFlagCorrectlyForPrivateProfile() =
         with(BgDataModel()) {
             setFlagsRule.enableFlags(Flags.FLAG_ENABLE_PRIVATE_SPACE)
-            val MAIN_HANDLE = UserHandle.of(0)
+            val MAIN_HANDLE = Process.myUserHandle()
             val mockUserHandles = arrayListOf<UserHandle>(MAIN_HANDLE)
             `when`(userCache.userProfiles).thenReturn(mockUserHandles)
             `when`(userManagerState?.isUserQuiet(MAIN_HANDLE)).thenReturn(true)
diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index d866a9f..4dba78a 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -146,7 +146,7 @@
         executeOnLauncher(l -> {
             // Ensure updates are not deferred so notification happens when apps pause.
             l.getAppsView().getAppsStore().disableDeferUpdates(DEFER_UPDATES_TEST);
-            l.getAppsView().getWorkManager().getWorkUtilityView().performClick();
+            l.getAppsView().getWorkManager().getWorkUtilityView().getWorkFAB().performClick();
         });
 
         waitForLauncherCondition("Work profile toggle OFF failed", launcher -> {
diff --git a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
index 2219003..a0f227e 100644
--- a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
@@ -21,6 +21,7 @@
 
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
+import android.os.Process;
 
 import androidx.annotation.Nullable;
 import androidx.test.InstrumentationRegistry;
@@ -109,8 +110,9 @@
     }
 
     public static String getLauncherCommand(ActivityInfo launcher) {
-        return "cmd package set-home-activity " +
-                new ComponentName(launcher.packageName, launcher.name).flattenToString();
+        return String.format("cmd package set-home-activity --user %d %s",
+                Process.myUserHandle().getIdentifier(),
+                new ComponentName(launcher.packageName, launcher.name).flattenToString());
     }
 
     /**
diff --git a/tests/tapl/com/android/launcher3/tapl/TestHelpers.java b/tests/tapl/com/android/launcher3/tapl/TestHelpers.java
index 7f6062f..b43bfcf 100644
--- a/tests/tapl/com/android/launcher3/tapl/TestHelpers.java
+++ b/tests/tapl/com/android/launcher3/tapl/TestHelpers.java
@@ -33,10 +33,9 @@
 import androidx.test.uiautomator.SearchCondition;
 import androidx.test.uiautomator.UiDevice;
 
-import org.junit.Assert;
-
 import java.util.Date;
 import java.util.List;
+import java.util.Objects;
 
 public class TestHelpers {
 
@@ -113,8 +112,8 @@
     }
 
     private static String checkCrash(Context context, String label, long startTime) {
-        DropBoxManager dropbox = (DropBoxManager) context.getSystemService(Context.DROPBOX_SERVICE);
-        Assert.assertNotNull("Unable access the DropBoxManager service", dropbox);
+        DropBoxManager dropbox = Objects.requireNonNull(
+                context.getSystemService(DropBoxManager.class));
 
         long timestamp = startTime;
         DropBoxManager.Entry entry;