diff --git a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
similarity index 78%
rename from tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
rename to tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 33ffd1d..d866a9f 100644
--- a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -17,9 +17,7 @@
 
 import static com.android.launcher3.LauncherPrefs.WORK_EDU_STEP;
 import static com.android.launcher3.LauncherState.ALL_APPS;
-import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
-import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
 import static com.android.launcher3.util.TestUtil.installDummyAppForUser;
 import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
 import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
@@ -43,36 +41,39 @@
 import com.android.launcher3.allapps.WorkEduCard;
 import com.android.launcher3.allapps.WorkPausedCard;
 import com.android.launcher3.allapps.WorkProfileManager;
-import com.android.launcher3.tapl.LauncherInstrumentation;
+import com.android.launcher3.util.BaseLauncherActivityTest;
 import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.rule.ScreenRecordRule;
+import com.android.launcher3.util.rule.TestStabilityRule;
 import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.io.IOException;
-import java.util.Objects;
 import java.util.function.Predicate;
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplWorkProfileTest extends AbstractLauncherUiTest<Launcher> {
+public class WorkProfileTest extends BaseLauncherActivityTest<Launcher> {
 
     private static final int WORK_PAGE = ActivityAllAppsContainerView.AdapterHolder.WORK;
+    public static final int WAIT_TIME_MS = 30000;
+
+    @Rule
+    public ScreenRecordRule mScreenRecordRule = new ScreenRecordRule();
+    @Rule
+    public TestStabilityRule mTestStabilityRule = new TestStabilityRule();
 
     private int mProfileUserId;
     private boolean mWorkProfileSetupSuccessful;
-    private final String TAG = "WorkProfileTest";
+    private static final String TAG = "WorkProfileTest";
 
     @Before
-    @Override
     public void setUp() throws Exception {
-        super.setUp();
-        String output =
-                mDevice.executeShellCommand(
-                        "pm create-user --profileOf 0 --managed TestProfile");
+        String output = executeShellCommand("pm create-user --profileOf 0 --managed TestProfile");
         updateWorkProfileSetupSuccessful("pm create-user", output);
 
         String[] tokens = output.split("\\s+");
@@ -88,36 +89,15 @@
             return; // no need to setup launcher since all tests will skip.
         }
 
-        mDevice.pressHome();
-        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForStateTransitionToEnd("Launcher internal state didn't switch to Normal",
-                () -> NORMAL);
-        waitForResumed("Launcher internal state is still Background");
-        mLauncher.getWorkspace().switchToAllApps();
-        waitForStateTransitionToEnd("Launcher internal state didn't switch to All Apps",
-                () -> ALL_APPS);
+        loadLauncherSync();
+        goToState(ALL_APPS);
+        waitForState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
     }
 
     @After
     public void removeWorkProfile() throws Exception {
-        executeOnLauncherInTearDown(launcher -> {
-            if (launcher.getAppsView() == null) {
-                return;
-            }
-            launcher.getAppsView().getAppsStore().disableDeferUpdates(DEFER_UPDATES_TEST);
-        });
         TestUtil.uninstallDummyApp();
-
-        mLauncher.runToState(
-                () -> {
-                    try {
-                        mDevice.executeShellCommand("pm remove-user --wait " + mProfileUserId);
-                    } catch (IOException e) {
-                        throw new RuntimeException(e);
-                    }
-                },
-                NORMAL_STATE_ORDINAL,
-                "executing pm 'remove-user' command");
+        executeShellCommand("pm remove-user --wait " + mProfileUserId);
     }
 
     private void waitForWorkTabSetup() {
@@ -127,7 +107,7 @@
                 return true;
             }
             return false;
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
     }
 
     @Test
@@ -137,10 +117,10 @@
         waitForWorkTabSetup();
         waitForLauncherCondition("Personal tab is missing",
                 launcher -> launcher.getAppsView().isPersonalTabVisible(),
-                LauncherInstrumentation.WAIT_TIME_MS);
+                WAIT_TIME_MS);
         waitForLauncherCondition("Work tab is missing",
                 launcher -> launcher.getAppsView().isWorkTabVisible(),
-                LauncherInstrumentation.WAIT_TIME_MS);
+                WAIT_TIME_MS);
     }
 
     // Staging; will be promoted to presubmit if stable
@@ -156,12 +136,11 @@
 
         WorkProfileManager manager = getFromLauncher(l -> l.getAppsView().getWorkManager());
 
-
         waitForLauncherCondition("work profile initial state check failed", launcher ->
                         manager.getWorkUtilityView() != null
                                 && manager.getCurrentState() == WorkProfileManager.STATE_ENABLED
                                 && manager.getWorkUtilityView().isEnabled(),
-                LauncherInstrumentation.WAIT_TIME_MS);
+                WAIT_TIME_MS);
 
         //start work profile toggle OFF test
         executeOnLauncher(l -> {
@@ -173,7 +152,7 @@
         waitForLauncherCondition("Work profile toggle OFF failed", launcher -> {
             manager.reset(); // pulls current state from system
             return manager.getCurrentState() == WorkProfileManager.STATE_DISABLED;
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
 
         waitForWorkCard("Work paused card not shown", view -> view instanceof WorkPausedCard);
 
@@ -188,7 +167,7 @@
         waitForLauncherCondition("Work profile toggle ON failed", launcher -> {
             manager.reset(); // pulls current state from system
             return manager.getCurrentState() == WorkProfileManager.STATE_ENABLED;
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
 
     }
 
@@ -215,7 +194,7 @@
             } finally {
                 l.getAppsView().getAppsStore().enableDeferUpdates(DEFER_UPDATES_TEST);
             }
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
     }
 
     private void updateWorkProfileSetupSuccessful(String cli, String output) {
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
deleted file mode 100644
index d653317..0000000
--- a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
+++ /dev/null
@@ -1,177 +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.ui.workspace;
-
-import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
-
-import android.content.ContentProviderClient;
-import android.content.ContentResolver;
-import android.content.ContentValues;
-import android.net.Uri;
-import android.view.View;
-import android.view.ViewGroup;
-
-import androidx.test.filters.LargeTest;
-
-import com.android.launcher3.BubbleTextView;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.icons.ThemedIconDrawable;
-import com.android.launcher3.tapl.HomeAllApps;
-import com.android.launcher3.tapl.HomeAppIcon;
-import com.android.launcher3.tapl.HomeAppIconMenuItem;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
-import com.android.launcher3.util.Executors;
-
-import org.junit.Test;
-
-import java.util.ArrayDeque;
-import java.util.Queue;
-
-/**
- * Tests for theme icon support in Launcher
- *
- * Note running these tests will clear the workspace on the device.
- */
-@LargeTest
-public class TaplThemeIconsTest extends AbstractLauncherUiTest<Launcher> {
-
-    private static final String APP_NAME = "IconThemedActivity";
-    private static final String SHORTCUT_NAME = "Shortcut 1";
-
-    @Test
-    public void testIconWithoutTheme() throws Exception {
-        setThemeEnabled(false);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(APP_NAME);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getAppsView(), false));
-            icon.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), false));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    @Test
-    public void testShortcutIconWithoutTheme() throws Exception {
-        setThemeEnabled(false);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(TEST_APP_NAME);
-            HomeAppIconMenuItem shortcutItem =
-                    (HomeAppIconMenuItem) icon.openDeepShortcutMenu().getMenuItem(SHORTCUT_NAME);
-            shortcutItem.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), false));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    @Test
-    public void testIconWithTheme() throws Exception {
-        setThemeEnabled(true);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(APP_NAME);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getAppsView(), false));
-            icon.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), true));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    @Test
-    public void testShortcutIconWithTheme() throws Exception {
-        setThemeEnabled(true);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(TEST_APP_NAME);
-            HomeAppIconMenuItem shortcutItem =
-                    (HomeAppIconMenuItem) icon.openDeepShortcutMenu().getMenuItem(SHORTCUT_NAME);
-            shortcutItem.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), true));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    private void verifyIconTheme(String title, ViewGroup parent, boolean isThemed) {
-        // Wait for Launcher model to be completed
-        try {
-            Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-
-        // Find the app icon
-        Queue<View> viewQueue = new ArrayDeque<>();
-        viewQueue.add(parent);
-        BubbleTextView icon = null;
-        while (!viewQueue.isEmpty()) {
-            View view = viewQueue.poll();
-            if (view instanceof ViewGroup) {
-                parent = (ViewGroup) view;
-                for (int i = parent.getChildCount() - 1; i >= 0; i--) {
-                    viewQueue.add(parent.getChildAt(i));
-                }
-            } else if (view instanceof BubbleTextView btv) {
-                if (title.equals(btv.getContentDescription().toString())) {
-                    icon = btv;
-                    break;
-                }
-            }
-        }
-
-        assertNotNull(icon.getIcon());
-        assertEquals(isThemed, icon.getIcon() instanceof ThemedIconDrawable);
-    }
-
-    private void setThemeEnabled(boolean isEnabled) throws Exception {
-        Uri uri = new Uri.Builder()
-                .scheme(ContentResolver.SCHEME_CONTENT)
-                .authority(mTargetPackage + ".grid_control")
-                .appendPath("set_icon_themed")
-                .build();
-        ContentValues values = new ContentValues();
-        values.put("boolean_value", isEnabled);
-        try (ContentProviderClient client = mTargetContext.getContentResolver()
-                .acquireContentProviderClient(uri)) {
-            int result = client.update(uri, values, null);
-            assertTrue(result > 0);
-        }
-    }
-}
diff --git a/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
new file mode 100644
index 0000000..cfc0a6b
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
@@ -0,0 +1,203 @@
+/*
+ * 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.ui.workspace;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.AbstractFloatingView.TYPE_ACTION_POPUP;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import android.content.ContentProviderClient;
+import android.content.ContentResolver;
+import android.content.ContentValues;
+import android.net.Uri;
+import android.view.View;
+import android.view.ViewGroup;
+
+import androidx.test.filters.LargeTest;
+import androidx.test.uiautomator.UiDevice;
+
+import com.android.launcher3.AbstractFloatingView;
+import com.android.launcher3.BubbleTextView;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.LauncherState;
+import com.android.launcher3.allapps.AllAppsRecyclerView;
+import com.android.launcher3.celllayout.FavoriteItemsTransaction;
+import com.android.launcher3.icons.ThemedIconDrawable;
+import com.android.launcher3.popup.ArrowPopup;
+import com.android.launcher3.util.BaseLauncherActivityTest;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.TestUtil;
+
+import org.junit.Test;
+
+import java.util.ArrayDeque;
+import java.util.Queue;
+
+/**
+ * Tests for theme icon support in Launcher
+ *
+ * Note running these tests will clear the workspace on the device.
+ */
+@LargeTest
+public class ThemeIconsTest extends BaseLauncherActivityTest<Launcher> {
+
+    private static final String APP_NAME = "IconThemedActivity";
+    private static final String SHORTCUT_NAME = "Shortcut 1";
+
+    @Test
+    public void testIconWithoutTheme() throws Exception {
+        setThemeEnabled(false);
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(APP_NAME);
+        BubbleTextView btv = getFromLauncher(
+                l -> verifyIconTheme(APP_NAME, l.getAppsView(), false));
+        addToWorkspace(btv);
+        executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), false));
+    }
+
+    @Test
+    public void testShortcutIconWithoutTheme() throws Exception {
+        setThemeEnabled(false);
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(TEST_APP_NAME);
+        BubbleTextView btv = getFromLauncher(l -> findBtv(TEST_APP_NAME, l.getAppsView()));
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, btv::performLongClick);
+
+        BubbleTextView menuItem = getOnceNotNull("Popup menu not open", l ->
+                (AbstractFloatingView.getOpenView(l, TYPE_ACTION_POPUP) instanceof ArrowPopup ap)
+                        ? findBtv(SHORTCUT_NAME, ap) : null);
+        addToWorkspace(menuItem);
+        executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), false));
+    }
+
+    @Test
+    public void testIconWithTheme() throws Exception {
+        setThemeEnabled(true);
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(APP_NAME);
+        BubbleTextView btv = getFromLauncher(l ->
+                verifyIconTheme(APP_NAME, l.getAppsView(), false));
+        addToWorkspace(btv);
+        executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), true));
+    }
+
+    @Test
+    public void testShortcutIconWithTheme() throws Exception {
+        setThemeEnabled(true);
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(TEST_APP_NAME);
+        BubbleTextView btv = getFromLauncher(l -> findBtv(TEST_APP_NAME, l.getAppsView()));
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, btv::performLongClick);
+
+        BubbleTextView menuItem = getOnceNotNull("Popup menu not open", l ->
+                (AbstractFloatingView.getOpenView(l, TYPE_ACTION_POPUP) instanceof ArrowPopup ap)
+                        ? findBtv(SHORTCUT_NAME, ap) : null);
+        addToWorkspace(menuItem);
+        executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), true));
+    }
+
+    private BubbleTextView findBtv(String title, ViewGroup parent) {
+        // Wait for Launcher model to be completed
+        try {
+            Executors.MODEL_EXECUTOR.submit(() -> { }).get();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        // Find the app icon
+        Queue<View> viewQueue = new ArrayDeque<>();
+        viewQueue.add(parent);
+        BubbleTextView icon = null;
+        while (!viewQueue.isEmpty()) {
+            View view = viewQueue.poll();
+            if (view instanceof ViewGroup) {
+                parent = (ViewGroup) view;
+                for (int i = parent.getChildCount() - 1; i >= 0; i--) {
+                    viewQueue.add(parent.getChildAt(i));
+                }
+            } else if (view instanceof BubbleTextView btv) {
+                if (btv.getContentDescription() != null
+                        && title.equals(btv.getContentDescription().toString())) {
+                    icon = btv;
+                    break;
+                }
+            }
+        }
+        return icon;
+    }
+
+    private BubbleTextView verifyIconTheme(String title, ViewGroup parent, boolean isThemed) {
+        BubbleTextView icon = findBtv(title, parent);
+        assertNotNull(icon.getIcon());
+        assertEquals(isThemed, icon.getIcon() instanceof ThemedIconDrawable);
+        return icon;
+    }
+
+    private void setThemeEnabled(boolean isEnabled) throws Exception {
+        Uri uri = new Uri.Builder()
+                .scheme(ContentResolver.SCHEME_CONTENT)
+                .authority(targetContext().getPackageName() + ".grid_control")
+                .appendPath("set_icon_themed")
+                .build();
+        ContentValues values = new ContentValues();
+        values.put("boolean_value", isEnabled);
+        try (ContentProviderClient client = targetContext().getContentResolver()
+                .acquireContentProviderClient(uri)) {
+            int result = client.update(uri, values, null);
+            assertTrue(result > 0);
+        }
+    }
+
+    private void scrollToAppIcon(String appName) {
+        executeOnLauncher(l -> {
+            l.hideKeyboard();
+            AllAppsRecyclerView rv = l.getAppsView().getActiveRecyclerView();
+            int pos = rv.getApps().getAdapterItems().indexOf(rv.getApps().getAdapterItems().stream()
+                    .filter(i -> i.itemInfo != null && appName.equals(i.itemInfo.title.toString()))
+                    .findFirst()
+                    .get());
+            rv.getLayoutManager().scrollToPosition(pos);
+        });
+    }
+
+    private void addToWorkspace(View btv) {
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, () ->
+                btv.getAccessibilityDelegate().performAccessibilityAction(
+                        btv, com.android.launcher3.R.id.action_add_to_workspace, null));
+        UiDevice.getInstance(getInstrumentation()).waitForIdle();
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt b/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
index bacce40..476e497 100644
--- a/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
+++ b/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
@@ -69,8 +69,9 @@
 
     protected fun targetContext(): Context = getInstrumentation().targetContext
 
-    protected fun goToState(state: LauncherState) = executeOnLauncher {
-        it.stateManager.goToState(state, 0)
+    protected fun goToState(state: LauncherState) {
+        executeOnLauncher { it.stateManager.goToState(state, 0) }
+        UiDevice.getInstance(getInstrumentation()).waitForIdle()
     }
 
     protected fun executeOnLauncher(f: ActivityAction<LAUNCHER_TYPE>) = scenario.onActivity(f)
@@ -92,6 +93,12 @@
         condition: Function<LAUNCHER_TYPE, Boolean>,
     ) = atMost(message, { getFromLauncher(condition)!! })
 
+    protected fun waitForLauncherCondition(
+        message: String,
+        condition: Function<LAUNCHER_TYPE, Boolean>,
+        timeout: Long,
+    ) = atMost(message, { getFromLauncher(condition)!! }, null, timeout)
+
     protected fun <T> getOnceNotNull(message: String, f: Function<LAUNCHER_TYPE, T?>): T? {
         var output: T? = null
         atMost(