diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 96622cafbe..fa7e8e93bb 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1123,6 +1123,16 @@ public final class LauncherInstrumentation {
         return object;
     }
 
+    @Nullable
+    UiObject2 findObjectInContainer(UiObject2 container, String resName) {
+        try {
+            return container.findObject(getLauncherObjectSelector(resName));
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
+    }
+
     @Nullable
     UiObject2 findObjectInContainer(UiObject2 container, BySelector selector) {
         try {
diff --git a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
index 0752e1e91b..adc993df94 100644
--- a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
+++ b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
@@ -182,6 +182,6 @@ public final class OverviewTask {
     }
 
     boolean isTaskSplit() {
-        return mTask.hasObject(By.res("bottomright_snapshot"));
+        return mLauncher.findObjectInContainer(mTask.getParent(), "bottomright_snapshot") != null;
     }
 }