diff --git a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
index 6a48b77..44c23ba 100644
--- a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
+++ b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
@@ -23,9 +23,9 @@
 import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.Until;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.tapl.LaunchedAppState;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.uioverrides.QuickstepLauncher;
 import com.android.quickstep.views.RecentsView;
 
 import org.junit.rules.RuleChain;
@@ -34,7 +34,7 @@
 /**
  * Base class for all instrumentation tests that deal with Quickstep.
  */
-public abstract class AbstractQuickStepTest extends AbstractLauncherUiTest {
+public abstract class AbstractQuickStepTest extends AbstractLauncherUiTest<QuickstepLauncher> {
     public static final boolean ENABLE_SHELL_TRANSITIONS =
             SystemProperties.getBoolean("persist.wm.debug.shell_transit", true);
     @Override
@@ -46,7 +46,7 @@
     }
 
     @Override
-    protected void onLauncherActivityClose(Launcher launcher) {
+    protected void onLauncherActivityClose(QuickstepLauncher launcher) {
         RecentsView recentsView = launcher.getOverviewPanel();
         if (recentsView != null) {
             recentsView.finishRecentsAnimation(false /* toRecents */, null);
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index b90839d..7dabbca 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -25,6 +25,7 @@
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
 import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.quickstep.views.LauncherRecentsView
 import com.android.quickstep.views.TaskView
@@ -45,7 +46,7 @@
 
     @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
 
-    private val launcher: Launcher = mock()
+    private val launcher: QuickstepLauncher = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
diff --git a/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java b/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java
index 3f3b9ed..fa10b61 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java
@@ -27,6 +27,7 @@
 import com.android.launcher3.tapl.OverviewTask.OverviewSplitTask;
 import com.android.launcher3.tapl.OverviewTaskMenu;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.uioverrides.QuickstepLauncher;
 import com.android.launcher3.util.rule.TestStabilityRule;
 
 import org.junit.Test;
@@ -35,7 +36,7 @@
  * This test run in both Out of process (Oop) and in-process (Ipc).
  * Tests the app Icon in overview.
  */
-public class TaplOverviewIconTest extends AbstractLauncherUiTest {
+public class TaplOverviewIconTest extends AbstractLauncherUiTest<QuickstepLauncher> {
 
     private static final String CALCULATOR_APP_PACKAGE =
             resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR);
diff --git a/quickstep/tests/src/com/android/quickstep/TaskViewTest.java b/quickstep/tests/src/com/android/quickstep/TaskViewTest.java
index d744194..8eec903 100644
--- a/quickstep/tests/src/com/android/quickstep/TaskViewTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaskViewTest.java
@@ -35,6 +35,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.statemanager.StatefulActivity;
+import com.android.launcher3.uioverrides.QuickstepLauncher;
 import com.android.quickstep.util.BorderAnimator;
 import com.android.quickstep.views.TaskView;
 
@@ -47,7 +48,7 @@
 public class TaskViewTest {
 
     @Mock
-    private StatefulActivity mContext;
+    private QuickstepLauncher mContext;
     @Mock
     private Resources mResource;
     @Mock
diff --git a/tests/src/com/android/launcher3/LauncherIntentTest.java b/tests/src/com/android/launcher3/LauncherIntentTest.java
index d1110c3..aeeb42a 100644
--- a/tests/src/com/android/launcher3/LauncherIntentTest.java
+++ b/tests/src/com/android/launcher3/LauncherIntentTest.java
@@ -34,7 +34,7 @@
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class LauncherIntentTest extends AbstractLauncherUiTest {
+public class LauncherIntentTest extends AbstractLauncherUiTest<Launcher> {
 
     public final Intent allAppsIntent = new Intent(Intent.ACTION_ALL_APPS);
 
diff --git a/tests/src/com/android/launcher3/allapps/TaplAllAppsIconsWorkingTest.java b/tests/src/com/android/launcher3/allapps/TaplAllAppsIconsWorkingTest.java
index 10e9f8a..0b233e5 100644
--- a/tests/src/com/android/launcher3/allapps/TaplAllAppsIconsWorkingTest.java
+++ b/tests/src/com/android/launcher3/allapps/TaplAllAppsIconsWorkingTest.java
@@ -18,6 +18,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.tapl.AppIcon;
 import com.android.launcher3.tapl.HomeAllApps;
@@ -30,7 +31,7 @@
  * The test runs in Out of process (Oop) and in process.
  * Makes sure the basic behaviors of Icons on AllApps are working.
  */
-public class TaplAllAppsIconsWorkingTest extends AbstractLauncherUiTest {
+public class TaplAllAppsIconsWorkingTest extends AbstractLauncherUiTest<Launcher> {
 
     /**
      * Makes sure we can launch an icon from All apps
diff --git a/tests/src/com/android/launcher3/allapps/TaplKeyboardFocusTest.java b/tests/src/com/android/launcher3/allapps/TaplKeyboardFocusTest.java
index 0360470..20684eb 100644
--- a/tests/src/com/android/launcher3/allapps/TaplKeyboardFocusTest.java
+++ b/tests/src/com/android/launcher3/allapps/TaplKeyboardFocusTest.java
@@ -26,6 +26,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.tapl.HomeAllApps;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
@@ -37,7 +38,7 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
-public class TaplKeyboardFocusTest extends AbstractLauncherUiTest {
+public class TaplKeyboardFocusTest extends AbstractLauncherUiTest<Launcher> {
 
     @Test
     public void testAllAppsFocusApp() {
diff --git a/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllAppsTest.java b/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllAppsTest.java
index e462c4f..f3f6fa5 100644
--- a/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllAppsTest.java
+++ b/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllAppsTest.java
@@ -28,6 +28,7 @@
 import androidx.test.filters.FlakyTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.tapl.AllApps;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
@@ -40,7 +41,7 @@
  * Test that we can open and close the all apps in multiple situations.
  * The test runs in Out of process (Oop) and in process.
  */
-public class TaplOpenCloseAllAppsTest extends AbstractLauncherUiTest {
+public class TaplOpenCloseAllAppsTest extends AbstractLauncherUiTest<Launcher> {
 
     public static final String READ_DEVICE_CONFIG_PERMISSION =
             "android.permission.READ_DEVICE_CONFIG";
diff --git a/tests/src/com/android/launcher3/appiconmenu/TaplAppIconMenuTest.java b/tests/src/com/android/launcher3/appiconmenu/TaplAppIconMenuTest.java
index c69d1b9..55d3252 100644
--- a/tests/src/com/android/launcher3/appiconmenu/TaplAppIconMenuTest.java
+++ b/tests/src/com/android/launcher3/appiconmenu/TaplAppIconMenuTest.java
@@ -38,7 +38,7 @@
  * Tests the AppIconMenu (the menu that appears when you long press an app icon) and also make sure
  * we can launch a shortcut from it.
  */
-public class TaplAppIconMenuTest extends AbstractLauncherUiTest {
+public class TaplAppIconMenuTest extends AbstractLauncherUiTest<Launcher> {
 
     private boolean isOptionsPopupVisible(Launcher launcher) {
         final ArrowPopup<?> popup = launcher.getOptionsPopup();
diff --git a/tests/src/com/android/launcher3/celllayout/TaplReorderWidgetsTest.java b/tests/src/com/android/launcher3/celllayout/TaplReorderWidgetsTest.java
index 1b03645..389d027 100644
--- a/tests/src/com/android/launcher3/celllayout/TaplReorderWidgetsTest.java
+++ b/tests/src/com/android/launcher3/celllayout/TaplReorderWidgetsTest.java
@@ -30,6 +30,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.MultipageCellLayout;
 import com.android.launcher3.celllayout.board.CellLayoutBoard;
@@ -59,7 +60,7 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
-public class TaplReorderWidgetsTest extends AbstractLauncherUiTest {
+public class TaplReorderWidgetsTest extends AbstractLauncherUiTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
diff --git a/tests/src/com/android/launcher3/compat/TaplPromiseIconUiTest.java b/tests/src/com/android/launcher3/compat/TaplPromiseIconUiTest.java
index 7ed0fa5..1500538 100644
--- a/tests/src/com/android/launcher3/compat/TaplPromiseIconUiTest.java
+++ b/tests/src/com/android/launcher3/compat/TaplPromiseIconUiTest.java
@@ -52,7 +52,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplPromiseIconUiTest extends AbstractLauncherUiTest {
+public class TaplPromiseIconUiTest extends AbstractLauncherUiTest<Launcher> {
 
     @Rule
     public final CheckFlagsRule mCheckFlagsRule =
diff --git a/tests/src/com/android/launcher3/dragging/TaplDragTest.java b/tests/src/com/android/launcher3/dragging/TaplDragTest.java
index bf1ba39..d64d096 100644
--- a/tests/src/com/android/launcher3/dragging/TaplDragTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplDragTest.java
@@ -33,6 +33,7 @@
 import android.platform.test.annotations.PlatinumTest;
 import android.util.Log;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.tapl.Folder;
 import com.android.launcher3.tapl.FolderIcon;
 import com.android.launcher3.tapl.HomeAllApps;
@@ -57,7 +58,7 @@
  *    * Can drag an icon from AllApps into the workspace
  *    * Can drag an icon on the Workspace to other positions of the Workspace.
  */
-public class TaplDragTest extends AbstractLauncherUiTest {
+public class TaplDragTest extends AbstractLauncherUiTest<Launcher> {
 
     /**
      * Adds two icons to the Workspace and combines them into a folder, then makes sure the icons
diff --git a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
index e21309f..df71ec0 100644
--- a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
@@ -17,7 +17,6 @@
 
 import static com.android.launcher3.testing.shared.TestProtocol.ICON_MISSING;
 import static com.android.launcher3.testing.shared.TestProtocol.UIOBJECT_STALE_ELEMENT;
-import static com.android.launcher3.ui.AbstractLauncherUiTest.initialize;
 import static com.android.launcher3.util.TestConstants.AppNames.DUMMY_APP_NAME;
 import static com.android.launcher3.util.TestConstants.AppNames.GMAIL_APP_NAME;
 import static com.android.launcher3.util.TestConstants.AppNames.MAPS_APP_NAME;
@@ -32,6 +31,7 @@
 import android.platform.test.annotations.PlatinumTest;
 import android.util.Log;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.tapl.HomeAllApps;
 import com.android.launcher3.tapl.HomeAppIcon;
 import com.android.launcher3.tapl.Workspace;
@@ -51,7 +51,7 @@
  * Test runs in Out of process (Oop) and In process (Ipc)
  * Test the behaviour of uninstalling and removing apps both from AllApps, Workspace and Hotseat.
  */
-public class TaplUninstallRemoveTest extends AbstractLauncherUiTest {
+public class TaplUninstallRemoveTest extends AbstractLauncherUiTest<Launcher> {
 
     /**
      * Deletes app both built-in and user-installed from the Workspace and makes sure it's no longer
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index d8635f9..7109083 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -95,7 +95,7 @@
 /**
  * Base class for all instrumentation tests providing various utility methods.
  */
-public abstract class AbstractLauncherUiTest {
+public abstract class AbstractLauncherUiTest<LAUNCHER_TYPE extends Launcher> {
 
     public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 5;
@@ -251,7 +251,7 @@
         final ViewCaptureRule viewCaptureRule = new ViewCaptureRule(
                 Launcher.ACTIVITY_TRACKER::getCreatedActivity);
         final RuleChain inner = RuleChain
-                .outerRule(new PortraitLandscapeRunner(this))
+                .outerRule(new PortraitLandscapeRunner<LAUNCHER_TYPE>(this))
                 .around(new FailureWatcher(mLauncher, viewCaptureRule::getViewCaptureData))
                 // .around(viewCaptureRule) // b/315482167
                 .around(new TestIsolationRule(mLauncher, true));
@@ -426,12 +426,12 @@
         }
     }
 
-    protected <T> T getFromLauncher(Function<Launcher, T> f) {
+    protected <T> T getFromLauncher(Function<LAUNCHER_TYPE, T> f) {
         if (!TestHelpers.isInLauncherProcess()) return null;
         return getOnUiThread(() -> f.apply(Launcher.ACTIVITY_TRACKER.getCreatedActivity()));
     }
 
-    protected void executeOnLauncher(Consumer<Launcher> f) {
+    protected void executeOnLauncher(Consumer<LAUNCHER_TYPE> f) {
         getFromLauncher(launcher -> {
             f.accept(launcher);
             return null;
@@ -441,7 +441,7 @@
     // Execute an action on Launcher, but forgive it when launcher is null.
     // Launcher can be null if teardown is happening after a failed setup step where launcher
     // activity failed to be created.
-    protected void executeOnLauncherInTearDown(Consumer<Launcher> f) {
+    protected void executeOnLauncherInTearDown(Consumer<LAUNCHER_TYPE> f) {
         executeOnLauncher(launcher -> {
             if (launcher != null) f.accept(launcher);
         });
@@ -469,20 +469,20 @@
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
     protected void waitForLauncherCondition(String
-            message, Function<Launcher, Boolean> condition) {
+            message, Function<LAUNCHER_TYPE, Boolean> condition) {
         waitForLauncherCondition(message, condition, DEFAULT_ACTIVITY_TIMEOUT);
     }
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
-    protected <T> T getOnceNotNull(String message, Function<Launcher, T> f) {
+    protected <O> O getOnceNotNull(String message, Function<LAUNCHER_TYPE, O> f) {
         return getOnceNotNull(message, f, DEFAULT_ACTIVITY_TIMEOUT);
     }
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
     protected void waitForLauncherCondition(
-            String message, Function<Launcher, Boolean> condition, long timeout) {
+            String message, Function<LAUNCHER_TYPE, Boolean> condition, long timeout) {
         verifyKeyguardInvisible();
         if (!TestHelpers.isInLauncherProcess()) return;
         Wait.atMost(message, () -> getFromLauncher(condition), timeout, mLauncher);
@@ -490,7 +490,7 @@
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
-    protected <T> T getOnceNotNull(String message, Function<Launcher, T> f, long timeout) {
+    protected <T> T getOnceNotNull(String message, Function<LAUNCHER_TYPE, T> f, long timeout) {
         if (!TestHelpers.isInLauncherProcess()) return null;
 
         final Object[] output = new Object[1];
@@ -506,7 +506,7 @@
     // flakiness.
     protected void waitForLauncherCondition(
             String message,
-            Runnable testThreadAction, Function<Launcher, Boolean> condition,
+            Runnable testThreadAction, Function<LAUNCHER_TYPE, Boolean> condition,
             long timeout) {
         if (!TestHelpers.isInLauncherProcess()) return;
         Wait.atMost(message, () -> {
@@ -648,7 +648,7 @@
                 "Launcher still active", launcher -> launcher == null, DEFAULT_UI_TIMEOUT);
     }
 
-    protected boolean isInLaunchedApp(Launcher launcher) {
+    protected boolean isInLaunchedApp(LAUNCHER_TYPE launcher) {
         return launcher == null || !launcher.hasBeenResumed();
     }
 
@@ -658,11 +658,11 @@
                 launcher -> launcher.getStateManager().getState() == state.get());
     }
 
-    protected int getAllAppsScroll(Launcher launcher) {
+    protected int getAllAppsScroll(LAUNCHER_TYPE launcher) {
         return launcher.getAppsView().getActiveRecyclerView().computeVerticalScrollOffset();
     }
 
-    protected void onLauncherActivityClose(Launcher launcher) {
+    protected void onLauncherActivityClose(LAUNCHER_TYPE launcher) {
     }
 
     protected HomeAppIcon createShortcutInCenterIfNotExist(String name) {
diff --git a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
index fabe3f7..c4e74f2 100644
--- a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
+++ b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
@@ -3,6 +3,7 @@
 import android.util.Log;
 import android.view.Surface;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.util.rule.FailureWatcher;
 
@@ -14,11 +15,12 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
+import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
-public class PortraitLandscapeRunner implements TestRule {
+public class PortraitLandscapeRunner<LAUNCHER_TYPE extends Launcher> implements TestRule {
     private static final String TAG = "PortraitLandscapeRunner";
-    private AbstractLauncherUiTest mTest;
+    private AbstractLauncherUiTest<LAUNCHER_TYPE> mTest;
 
     // Annotation for tests that need to be run in portrait and landscape modes.
     @Retention(RetentionPolicy.RUNTIME)
@@ -26,7 +28,7 @@
     public @interface PortraitLandscape {
     }
 
-    public PortraitLandscapeRunner(AbstractLauncherUiTest test) {
+    public PortraitLandscapeRunner(AbstractLauncherUiTest<LAUNCHER_TYPE> test) {
         mTest = test;
     }
 
@@ -47,7 +49,7 @@
 
                         mTest.mDevice.pressHome();
                         mTest.waitForLauncherCondition("Launcher activity wasn't created",
-                                launcher -> launcher != null,
+                                Objects::nonNull,
                                 TimeUnit.SECONDS.toMillis(20));
 
                         mTest.executeOnLauncher(launcher ->
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
index 6780ed5..342eedf 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
@@ -22,13 +22,14 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
+import com.android.launcher3.Launcher;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplTestsLauncher3Test extends AbstractLauncherUiTest {
+public class TaplTestsLauncher3Test extends AbstractLauncherUiTest<Launcher> {
 
     @ScreenRecord // b/322823478
     @Test
diff --git a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
index eb1f49f..cccda88 100644
--- a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
@@ -35,6 +35,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.R;
 import com.android.launcher3.allapps.ActivityAllAppsContainerView;
@@ -58,7 +59,7 @@
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplWorkProfileTest extends AbstractLauncherUiTest {
+public class TaplWorkProfileTest extends AbstractLauncherUiTest<Launcher> {
 
     private static final int WORK_PAGE = ActivityAllAppsContainerView.AdapterHolder.WORK;
 
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
index 1db6014..e6e02b4 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
@@ -52,7 +52,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplAddConfigWidgetTest extends AbstractLauncherUiTest {
+public class TaplAddConfigWidgetTest extends AbstractLauncherUiTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
index 6c24329..7845dec 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
@@ -25,6 +25,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.celllayout.FavoriteItemsTransaction;
 import com.android.launcher3.tapl.Widget;
 import com.android.launcher3.tapl.WidgetResizeFrame;
@@ -46,7 +47,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplAddWidgetTest extends AbstractLauncherUiTest {
+public class TaplAddWidgetTest extends AbstractLauncherUiTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
index e69cb4c..28d1faa 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
@@ -43,6 +43,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 
+import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherSettings;
@@ -76,7 +77,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplBindWidgetTest extends AbstractLauncherUiTest {
+public class TaplBindWidgetTest extends AbstractLauncherUiTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
index 90a6d22..74047f0 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
@@ -59,7 +59,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplRequestPinItemTest extends AbstractLauncherUiTest {
+public class TaplRequestPinItemTest extends AbstractLauncherUiTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java b/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
index 4acdddc..79724e1 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
@@ -43,7 +43,7 @@
  */
 @MediumTest
 @RunWith(AndroidJUnit4.class)
-public class TaplWidgetPickerTest extends AbstractLauncherUiTest {
+public class TaplWidgetPickerTest extends AbstractLauncherUiTest<Launcher> {
     // b/325377690 : To get the log printed where DOWN key event is getting lost from TAPL.
     @Rule public final TestRule mEnableInputTransportPublisherRule =
             createEnableInputTransportPublisherRule();
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
index c60ee8c..d653317 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
@@ -31,6 +31,7 @@
 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;
@@ -49,7 +50,7 @@
  * Note running these tests will clear the workspace on the device.
  */
 @LargeTest
-public class TaplThemeIconsTest extends AbstractLauncherUiTest {
+public class TaplThemeIconsTest extends AbstractLauncherUiTest<Launcher> {
 
     private static final String APP_NAME = "IconThemedActivity";
     private static final String SHORTCUT_NAME = "Shortcut 1";
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
index 913dfa2..2ce8eef 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
@@ -55,7 +55,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplTwoPanelWorkspaceTest extends AbstractLauncherUiTest {
+public class TaplTwoPanelWorkspaceTest extends AbstractLauncherUiTest<Launcher> {
 
     private AutoCloseable mLauncherLayout;
 
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplWorkspaceTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplWorkspaceTest.java
index 4edeb42..9afde0d 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplWorkspaceTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplWorkspaceTest.java
@@ -37,7 +37,7 @@
 /**
  * Test the basic interactions of the Workspace, adding pages, moving the pages and removing pages.
  */
-public class TaplWorkspaceTest extends AbstractLauncherUiTest {
+public class TaplWorkspaceTest extends AbstractLauncherUiTest<Launcher> {
 
     private AutoCloseable mLauncherLayout;