diff --git a/Android.bp b/Android.bp
index eca0778..61042f6 100644
--- a/Android.bp
+++ b/Android.bp
@@ -108,6 +108,7 @@
     ],
     srcs: [
         "tests/multivalentTests/tapl/**/*.java",
+        "tests/multivalentTests/tapl/**/*.kt",
     ],
     resource_dirs: [],
     manifest: "tests/multivalentTests/tapl/AndroidManifest.xml",
diff --git a/quickstep/res/layout/task.xml b/quickstep/res/layout/task.xml
index 823a86e..9d599c9 100644
--- a/quickstep/res/layout/task.xml
+++ b/quickstep/res/layout/task.xml
@@ -19,6 +19,7 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
     xmlns:launcher="http://schemas.android.com/apk/res-auto"
+    android:id="@+id/task"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:clipChildren="false"
diff --git a/quickstep/res/layout/task_desktop.xml b/quickstep/res/layout/task_desktop.xml
index 60827cd..3cafcfd 100644
--- a/quickstep/res/layout/task_desktop.xml
+++ b/quickstep/res/layout/task_desktop.xml
@@ -19,6 +19,7 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
     xmlns:launcher="http://schemas.android.com/apk/res-auto"
+    android:id="@+id/task"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:clipChildren="true"
diff --git a/quickstep/res/layout/task_grouped.xml b/quickstep/res/layout/task_grouped.xml
index d20afd3..e91e773 100644
--- a/quickstep/res/layout/task_grouped.xml
+++ b/quickstep/res/layout/task_grouped.xml
@@ -24,6 +24,7 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
     xmlns:launcher="http://schemas.android.com/apk/res-auto"
+    android:id="@+id/task"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:clipChildren="false"
diff --git a/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java b/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java
index 3f806d1..7c1b7f3 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplOverviewIconTest.java
@@ -15,7 +15,6 @@
  */
 package com.android.quickstep;
 
-import static com.android.launcher3.ui.AbstractLauncherUiTest.initialize;
 import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
 import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
@@ -25,6 +24,7 @@
 import android.content.Intent;
 import android.platform.test.annotations.PlatinumTest;
 
+import com.android.launcher3.tapl.OverviewTask.OverviewSplitTask;
 import com.android.launcher3.tapl.OverviewTaskMenu;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.rule.TestStabilityRule;
@@ -86,7 +86,8 @@
         taskMenu.touchOutsideTaskMenuToDismiss();
 
         OverviewTaskMenu splitMenu =
-                mLauncher.goHome().switchToOverview().getCurrentTask().tapSplitTaskMenu();
+                mLauncher.goHome().switchToOverview().getCurrentTask().tapMenu(
+                        OverviewSplitTask.SPLIT_BOTTOM_OR_RIGHT);
         assertTrue("App info item not appearing in expanded split task's menu.",
                 splitMenu.hasMenuItem("App info"));
         splitMenu.touchOutsideTaskMenuToDismiss();
diff --git a/tests/multivalentTests/tapl/com/android/launcher3/tapl/Background.java b/tests/multivalentTests/tapl/com/android/launcher3/tapl/Background.java
index 9f2ce22..988aa94 100644
--- a/tests/multivalentTests/tapl/com/android/launcher3/tapl/Background.java
+++ b/tests/multivalentTests/tapl/com/android/launcher3/tapl/Background.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.tapl;
 
+import static com.android.launcher3.tapl.BaseOverview.TASK_RES_ID;
 import static com.android.launcher3.tapl.OverviewTask.TASK_START_EVENT;
 import static com.android.launcher3.testing.shared.TestProtocol.OVERVIEW_STATE_ORDINAL;
 
@@ -116,10 +117,10 @@
                         // non-tablet overview, snapshots can be on either side of the swiped
                         // task, but we still check that they become visible after swiping and
                         // pausing.
-                        mLauncher.waitForOverviewObject("snapshot");
+                        mLauncher.waitForOverviewObject(TASK_RES_ID);
                         if (mLauncher.isTablet()) {
                             List<UiObject2> tasks = mLauncher.getDevice().findObjects(
-                                    mLauncher.getOverviewObjectSelector("snapshot"));
+                                    mLauncher.getOverviewObjectSelector(TASK_RES_ID));
                             final int centerX = mLauncher.getDevice().getDisplayWidth() / 2;
                             mLauncher.assertTrue(
                                     "All tasks not to the left of the swiped task",
diff --git a/tests/multivalentTests/tapl/com/android/launcher3/tapl/BaseOverview.java b/tests/multivalentTests/tapl/com/android/launcher3/tapl/BaseOverview.java
index 1bc489c..091d2fd 100644
--- a/tests/multivalentTests/tapl/com/android/launcher3/tapl/BaseOverview.java
+++ b/tests/multivalentTests/tapl/com/android/launcher3/tapl/BaseOverview.java
@@ -42,6 +42,7 @@
  * Common overview panel for both Launcher and fallback recents
  */
 public class BaseOverview extends LauncherInstrumentation.VisibleContainer {
+    protected static final String TASK_RES_ID = "task";
 
     private static final Pattern EVENT_ALT_ESC_DOWN = Pattern.compile(
             "Key event: KeyEvent.*?action=ACTION_DOWN.*?keyCode=KEYCODE_ESCAPE.*?metaState=0");
@@ -275,10 +276,10 @@
 
         // The widest, and most top-right task should be the current task
         UiObject2 currentTask = Collections.max(taskViews,
-                Comparator.comparingInt((UiObject2 t) -> t.getParent().getVisibleBounds().width())
-                        .thenComparingInt((UiObject2 t) -> t.getParent().getVisibleCenter().x)
+                Comparator.comparingInt((UiObject2 t) -> t.getVisibleBounds().width())
+                        .thenComparingInt((UiObject2 t) -> t.getVisibleCenter().x)
                         .thenComparing(Comparator.comparing(
-                                (UiObject2 t) -> t.getParent().getVisibleCenter().y).reversed()));
+                                (UiObject2 t) -> t.getVisibleCenter().y).reversed()));
         return new OverviewTask(mLauncher, currentTask, this);
     }
 
@@ -323,10 +324,11 @@
                 "want to get overview tasks")) {
             verifyActiveContainer();
             return mLauncher.getDevice().findObjects(
-                    mLauncher.getOverviewObjectSelector("snapshot"));
+                    mLauncher.getOverviewObjectSelector(TASK_RES_ID));
         }
     }
 
+
     int getTaskCount() {
         return getTasks().size();
     }
diff --git a/tests/multivalentTests/tapl/com/android/launcher3/tapl/OverviewTask.java b/tests/multivalentTests/tapl/com/android/launcher3/tapl/OverviewTask.java
index f383e99..afe5722 100644
--- a/tests/multivalentTests/tapl/com/android/launcher3/tapl/OverviewTask.java
+++ b/tests/multivalentTests/tapl/com/android/launcher3/tapl/OverviewTask.java
@@ -16,6 +16,10 @@
 
 package com.android.launcher3.tapl;
 
+import static com.android.launcher3.tapl.OverviewTask.OverviewSplitTask.DEFAULT;
+import static com.android.launcher3.tapl.OverviewTask.OverviewSplitTask.SPLIT_BOTTOM_OR_RIGHT;
+import static com.android.launcher3.tapl.OverviewTask.OverviewSplitTask.SPLIT_TOP_OR_LEFT;
+
 import android.graphics.Rect;
 
 import androidx.annotation.NonNull;
@@ -34,9 +38,6 @@
  */
 public final class OverviewTask {
     private static final String SYSTEMUI_PACKAGE = "com.android.systemui";
-    private static final String TASK_SNAPSHOT_1 = "snapshot";
-    private static final String TASK_SNAPSHOT_2 = "bottomright_snapshot";
-
     static final Pattern TASK_START_EVENT = Pattern.compile("startActivityFromRecentsAsync");
     static final Pattern SPLIT_SELECT_EVENT = Pattern.compile("enterSplitSelect");
     static final Pattern SPLIT_START_EVENT = Pattern.compile("launchSplitTasks");
@@ -64,15 +65,16 @@
             return getCombinedSplitTaskHeight();
         }
 
-        return mTask.getVisibleBounds().height();
+        UiObject2 taskSnapshot1 = findObjectInTask(DEFAULT.snapshotRes);
+        return taskSnapshot1.getVisibleBounds().height();
     }
 
     /**
      * Calculates the visible height for split tasks, containing 2 snapshot tiles and a divider.
      */
     private int getCombinedSplitTaskHeight() {
-        UiObject2 taskSnapshot1 = findObjectInTask(TASK_SNAPSHOT_1);
-        UiObject2 taskSnapshot2 = findObjectInTask(TASK_SNAPSHOT_2);
+        UiObject2 taskSnapshot1 = findObjectInTask(SPLIT_TOP_OR_LEFT.snapshotRes);
+        UiObject2 taskSnapshot2 = findObjectInTask(SPLIT_BOTTOM_OR_RIGHT.snapshotRes);
 
         // If the split task is partly off screen, taskSnapshot1 can be invisible.
         if (taskSnapshot1 == null) {
@@ -96,15 +98,16 @@
             return getCombinedSplitTaskWidth();
         }
 
-        return mTask.getVisibleBounds().width();
+        UiObject2 taskSnapshot1 = findObjectInTask(DEFAULT.snapshotRes);
+        return taskSnapshot1.getVisibleBounds().width();
     }
 
     /**
      * Calculates the visible width for split tasks, containing 2 snapshot tiles and a divider.
      */
     private int getCombinedSplitTaskWidth() {
-        UiObject2 taskSnapshot1 = findObjectInTask(TASK_SNAPSHOT_1);
-        UiObject2 taskSnapshot2 = findObjectInTask(TASK_SNAPSHOT_2);
+        UiObject2 taskSnapshot1 = findObjectInTask(SPLIT_TOP_OR_LEFT.snapshotRes);
+        UiObject2 taskSnapshot2 = findObjectInTask(SPLIT_BOTTOM_OR_RIGHT.snapshotRes);
 
         int left = Math.min(
                 taskSnapshot1.getVisibleBounds().left, taskSnapshot2.getVisibleBounds().left);
@@ -115,15 +118,15 @@
     }
 
     int getTaskCenterX() {
-        return mTask.getParent().getVisibleCenter().x;
+        return mTask.getVisibleCenter().x;
     }
 
     int getTaskCenterY() {
-        return mTask.getParent().getVisibleCenter().y;
+        return mTask.getVisibleCenter().y;
     }
 
     float getExactCenterX() {
-        return mTask.getParent().getVisibleBounds().exactCenterX();
+        return mTask.getVisibleBounds().exactCenterX();
     }
 
     UiObject2 getUiObject() {
@@ -225,11 +228,17 @@
     /** Taps the task menu. Returns the task menu object. */
     @NonNull
     public OverviewTaskMenu tapMenu() {
+        return tapMenu(DEFAULT);
+    }
+
+    /** Taps the task menu of the split task. Returns the split task's menu object. */
+    @NonNull
+    public OverviewTaskMenu tapMenu(OverviewSplitTask task) {
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "want to tap the task menu")) {
             mLauncher.clickLauncherObject(
-                    mLauncher.waitForObjectInContainer(mTask.getParent(), "icon"));
+                    mLauncher.waitForObjectInContainer(mTask, task.iconAppRes));
 
             try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
                     "tapped the task menu")) {
@@ -238,27 +247,31 @@
         }
     }
 
-    /** Taps the task menu of the split task. Returns the split task's menu object. */
-    @NonNull
-    public OverviewTaskMenu tapSplitTaskMenu() {
-        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
-             LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
-                     "want to tap the split task's menu")) {
-            mLauncher.clickLauncherObject(
-                    mLauncher.waitForObjectInContainer(mTask.getParent(), "bottomRight_icon"));
-
-            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
-                    "tapped the split task's menu")) {
-                return new OverviewTaskMenu(mLauncher);
-            }
-        }
-    }
-
     boolean isTaskSplit() {
-        return findObjectInTask(TASK_SNAPSHOT_2) != null;
+        return findObjectInTask(SPLIT_BOTTOM_OR_RIGHT.snapshotRes) != null;
     }
 
     private UiObject2 findObjectInTask(String resName) {
-        return mTask.getParent().findObject(mLauncher.getOverviewObjectSelector(resName));
+        return mTask.findObject(mLauncher.getOverviewObjectSelector(resName));
+    }
+
+    /**
+     * Enum used to specify  which task is retrieved when it is a split task.
+     */
+    public enum OverviewSplitTask {
+        // The main task when the task is not split.
+        DEFAULT("snapshot", "icon"),
+        // The first task in split task.
+        SPLIT_TOP_OR_LEFT("snapshot", "icon"),
+        // The second task in split task.
+        SPLIT_BOTTOM_OR_RIGHT("bottomright_snapshot", "bottomRight_icon");
+
+        public final String snapshotRes;
+        public final String iconAppRes;
+
+        OverviewSplitTask(String snapshotRes, String iconAppRes) {
+            this.snapshotRes = snapshotRes;
+            this.iconAppRes = iconAppRes;
+        }
     }
 }