Improving TAPL diagnostics

Bug: 128531133
Change-Id: If81861afccb665e2dc2293f2704d99818b3297c8
diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index 7ab91bc..b785007 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -60,47 +60,59 @@
      */
     @NonNull
     public AppIcon getAppIcon(String appName) {
-        final UiObject2 allAppsContainer = verifyActiveContainer();
-        final BySelector appIconSelector = AppIcon.getAppIconSelector(appName, mLauncher);
-        if (!hasClickableIcon(allAppsContainer, appIconSelector)) {
-            scrollBackToBeginning();
-            int attempts = 0;
-            while (!hasClickableIcon(allAppsContainer, appIconSelector) &&
-                    allAppsContainer.scroll(Direction.DOWN, 0.8f)) {
-                LauncherInstrumentation.assertTrue(
-                        "Exceeded max scroll attempts: " + MAX_SCROLL_ATTEMPTS,
-                        ++attempts <= MAX_SCROLL_ATTEMPTS);
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to get app icon on all apps")) {
+            final UiObject2 allAppsContainer = verifyActiveContainer();
+            final BySelector appIconSelector = AppIcon.getAppIconSelector(appName, mLauncher);
+            if (!hasClickableIcon(allAppsContainer, appIconSelector)) {
+                scrollBackToBeginning();
+                int attempts = 0;
+                try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer("scrolled")) {
+                    while (!hasClickableIcon(allAppsContainer, appIconSelector) &&
+                            allAppsContainer.scroll(Direction.DOWN, 0.8f)) {
+                        mLauncher.assertTrue(
+                                "Exceeded max scroll attempts: " + MAX_SCROLL_ATTEMPTS,
+                                ++attempts <= MAX_SCROLL_ATTEMPTS);
+                        verifyActiveContainer();
+                    }
+                }
                 verifyActiveContainer();
             }
-        }
-        verifyActiveContainer();
 
-        final UiObject2 appIcon = mLauncher.getObjectInContainer(allAppsContainer, appIconSelector);
-        ensureIconVisible(appIcon, allAppsContainer);
-        return new AppIcon(mLauncher, appIcon);
+            final UiObject2 appIcon = mLauncher.getObjectInContainer(allAppsContainer,
+                    appIconSelector);
+            ensureIconVisible(appIcon, allAppsContainer);
+            return new AppIcon(mLauncher, appIcon);
+        }
     }
 
     private void scrollBackToBeginning() {
-        final UiObject2 allAppsContainer = verifyActiveContainer();
-        final UiObject2 searchBox =
-                mLauncher.waitForObjectInContainer(allAppsContainer, "search_container_all_apps");
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to scroll back in all apps")) {
+            final UiObject2 allAppsContainer = verifyActiveContainer();
+            final UiObject2 searchBox =
+                    mLauncher.waitForObjectInContainer(allAppsContainer,
+                            "search_container_all_apps");
 
-        int attempts = 0;
-        allAppsContainer.setGestureMargins(0, searchBox.getVisibleBounds().bottom + 1, 0, 5);
+            int attempts = 0;
+            allAppsContainer.setGestureMargins(0, searchBox.getVisibleBounds().bottom + 1, 0, 5);
 
-        for (int scroll = getScroll(allAppsContainer);
-                scroll != 0;
-                scroll = getScroll(allAppsContainer)) {
-            LauncherInstrumentation.assertTrue("Negative scroll position", scroll > 0);
+            for (int scroll = getScroll(allAppsContainer);
+                    scroll != 0;
+                    scroll = getScroll(allAppsContainer)) {
+                mLauncher.assertTrue("Negative scroll position", scroll > 0);
 
-            LauncherInstrumentation.assertTrue(
-                    "Exceeded max scroll attempts: " + MAX_SCROLL_ATTEMPTS,
-                    ++attempts <= MAX_SCROLL_ATTEMPTS);
+                mLauncher.assertTrue(
+                        "Exceeded max scroll attempts: " + MAX_SCROLL_ATTEMPTS,
+                        ++attempts <= MAX_SCROLL_ATTEMPTS);
 
-            allAppsContainer.scroll(Direction.UP, 1);
+                allAppsContainer.scroll(Direction.UP, 1);
+            }
+
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer("scrolled up")) {
+                verifyActiveContainer();
+            }
         }
-
-        verifyActiveContainer();
     }
 
     private int getScroll(UiObject2 allAppsContainer) {
@@ -115,8 +127,11 @@
             // to reveal the app icon to have the MIN_INTERACT_SIZE
             final float pct = Math.max(((float) (MIN_INTERACT_SIZE - appHeight)) / mHeight, 0.2f);
             allAppsContainer.scroll(Direction.DOWN, pct);
-            mLauncher.waitForIdle();
-            verifyActiveContainer();
+            try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                    "scrolled an icon in all apps to make it visible - and then")) {
+                mLauncher.waitForIdle();
+                verifyActiveContainer();
+            }
         }
     }
 
@@ -124,21 +139,29 @@
      * Flings forward (down) and waits the fling's end.
      */
     public void flingForward() {
-        final UiObject2 allAppsContainer = verifyActiveContainer();
-        // Start the gesture in the center to avoid starting at elements near the top.
-        allAppsContainer.setGestureMargins(0, 0, 0, mHeight / 2);
-        allAppsContainer.fling(Direction.DOWN, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        verifyActiveContainer();
+        try(LauncherInstrumentation.Closable c =
+                    mLauncher.addContextLayer("want to fling forward in all apps")) {
+            final UiObject2 allAppsContainer = verifyActiveContainer();
+            // Start the gesture in the center to avoid starting at elements near the top.
+            allAppsContainer.setGestureMargins(0, 0, 0, mHeight / 2);
+            allAppsContainer.fling(Direction.DOWN,
+                    (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            verifyActiveContainer();
+        }
     }
 
     /**
      * Flings backward (up) and waits the fling's end.
      */
     public void flingBackward() {
-        final UiObject2 allAppsContainer = verifyActiveContainer();
-        // Start the gesture in the center, for symmetry with forward.
-        allAppsContainer.setGestureMargins(0, mHeight / 2, 0, 0);
-        allAppsContainer.fling(Direction.UP, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        verifyActiveContainer();
+        try(LauncherInstrumentation.Closable c =
+                    mLauncher.addContextLayer("want to fling backward in all apps")) {
+            final UiObject2 allAppsContainer = verifyActiveContainer();
+            // Start the gesture in the center, for symmetry with forward.
+            allAppsContainer.setGestureMargins(0, mHeight / 2, 0, 0);
+            allAppsContainer.fling(Direction.UP,
+                    (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            verifyActiveContainer();
+        }
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/AllAppsFromOverview.java b/tests/tapl/com/android/launcher3/tapl/AllAppsFromOverview.java
index dcc51b5..c3b671b 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllAppsFromOverview.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllAppsFromOverview.java
@@ -42,20 +42,24 @@
      */
     @NonNull
     public Overview switchBackToOverview() {
-        final UiObject2 allAppsContainer = verifyActiveContainer();
-        // Swipe from the search box to the bottom.
-        final UiObject2 qsb = mLauncher.waitForObjectInContainer(
-                allAppsContainer, "search_container_all_apps");
-        final Point start = qsb.getVisibleCenter();
-        final int swipeHeight = mLauncher.getTestInfo(
-                TestProtocol.REQUEST_ALL_APPS_TO_OVERVIEW_SWIPE_HEIGHT).
-                getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to switch back from all apps to overview")) {
+            final UiObject2 allAppsContainer = verifyActiveContainer();
+            // Swipe from the search box to the bottom.
+            final UiObject2 qsb = mLauncher.waitForObjectInContainer(
+                    allAppsContainer, "search_container_all_apps");
+            final Point start = qsb.getVisibleCenter();
+            final int swipeHeight = mLauncher.getTestInfo(
+                    TestProtocol.REQUEST_ALL_APPS_TO_OVERVIEW_SWIPE_HEIGHT).
+                    getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
 
-        final int endY = start.y + swipeHeight + mLauncher.getTouchSlop();
-        LauncherInstrumentation.log("AllAppsFromOverview.switchBackToOverview before swipe");
-        mLauncher.swipe(start.x, start.y, start.x, endY, OVERVIEW_STATE_ORDINAL);
+            final int endY = start.y + swipeHeight + mLauncher.getTouchSlop();
+            LauncherInstrumentation.log("AllAppsFromOverview.switchBackToOverview before swipe");
+            mLauncher.swipe(start.x, start.y, start.x, endY, OVERVIEW_STATE_ORDINAL);
 
-        return new Overview(mLauncher);
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer("swiped down")) {
+                return new Overview(mLauncher);
+            }
+        }
     }
-
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Background.java b/tests/tapl/com/android/launcher3/tapl/Background.java
index ef509ed..26c0ca4 100644
--- a/tests/tapl/com/android/launcher3/tapl/Background.java
+++ b/tests/tapl/com/android/launcher3/tapl/Background.java
@@ -57,11 +57,14 @@
      */
     @NonNull
     public BaseOverview switchToOverview() {
-        verifyActiveContainer();
-        goToOverviewUnchecked(BACKGROUND_APP_STATE_ORDINAL);
-        assertTrue("Overview not visible", mLauncher.getDevice().wait(
-                Until.hasObject(By.pkg(getOverviewPackageName())), WAIT_TIME_MS));
-        return new BaseOverview(mLauncher);
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to switch from background to overview")) {
+            verifyActiveContainer();
+            goToOverviewUnchecked(BACKGROUND_APP_STATE_ORDINAL);
+            assertTrue("Overview not visible", mLauncher.getDevice().wait(
+                    Until.hasObject(By.pkg(getOverviewPackageName())), WAIT_TIME_MS));
+            return new BaseOverview(mLauncher);
+        }
     }
 
     protected void goToOverviewUnchecked(int expectedState) {
diff --git a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
index 234c3bf..4205188 100644
--- a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
+++ b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
@@ -44,42 +44,54 @@
      * Flings forward (left) and waits the fling's end.
      */
     public void flingForward() {
-        LauncherInstrumentation.log("Overview.flingForward before fling");
-        final UiObject2 overview = verifyActiveContainer();
-        final int margin = (int) (50 * mLauncher.getDisplayDensity()) + 1;
-        overview.setGestureMargins(margin, 0, 0, 0);
-        overview.fling(Direction.LEFT, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        mLauncher.waitForIdle();
-        verifyActiveContainer();
+        try (LauncherInstrumentation.Closable c =
+                     mLauncher.addContextLayer("want to fling forward in overview")) {
+            LauncherInstrumentation.log("Overview.flingForward before fling");
+            final UiObject2 overview = verifyActiveContainer();
+            final int margin = (int) (50 * mLauncher.getDisplayDensity()) + 1;
+            overview.setGestureMargins(margin, 0, 0, 0);
+            overview.fling(Direction.LEFT, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            mLauncher.waitForIdle();
+            verifyActiveContainer();
+        }
     }
 
     /**
      * Dismissed all tasks by scrolling to Clear-all button and pressing it.
      */
     public Workspace dismissAllTasks() {
-        final BySelector clearAllSelector = mLauncher.getLauncherObjectSelector("clear_all");
-        for (int i = 0;
-                i < FLINGS_FOR_DISMISS_LIMIT
-                        && !verifyActiveContainer().hasObject(clearAllSelector);
-                ++i) {
-            flingForward();
-        }
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "dismissing all tasks")) {
+            final BySelector clearAllSelector = mLauncher.getLauncherObjectSelector("clear_all");
+            for (int i = 0;
+                    i < FLINGS_FOR_DISMISS_LIMIT
+                            && !verifyActiveContainer().hasObject(clearAllSelector);
+                    ++i) {
+                flingForward();
+            }
 
-        mLauncher.getObjectInContainer(verifyActiveContainer(), clearAllSelector).click();
-        return new Workspace(mLauncher);
+            mLauncher.getObjectInContainer(verifyActiveContainer(), clearAllSelector).click();
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
+                    "dismissed all tasks")) {
+                return new Workspace(mLauncher);
+            }
+        }
     }
 
     /**
      * Flings backward (right) and waits the fling's end.
      */
     public void flingBackward() {
-        LauncherInstrumentation.log("Overview.flingBackward before fling");
-        final UiObject2 overview = verifyActiveContainer();
-        final int margin = (int) (50 * mLauncher.getDisplayDensity()) + 1;
-        overview.setGestureMargins(0, 0, margin, 0);
-        overview.fling(Direction.RIGHT, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        mLauncher.waitForIdle();
-        verifyActiveContainer();
+        try (LauncherInstrumentation.Closable c =
+                     mLauncher.addContextLayer("want to fling backward in overview")) {
+            LauncherInstrumentation.log("Overview.flingBackward before fling");
+            final UiObject2 overview = verifyActiveContainer();
+            final int margin = (int) (50 * mLauncher.getDisplayDensity()) + 1;
+            overview.setGestureMargins(0, 0, margin, 0);
+            overview.fling(Direction.RIGHT, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            mLauncher.waitForIdle();
+            verifyActiveContainer();
+        }
     }
 
     /**
@@ -89,18 +101,21 @@
      */
     @NonNull
     public OverviewTask getCurrentTask() {
-        verifyActiveContainer();
-        final List<UiObject2> taskViews = mLauncher.getDevice().findObjects(
-                mLauncher.getLauncherObjectSelector("snapshot"));
-        LauncherInstrumentation.assertNotEquals("Unable to find a task", 0, taskViews.size());
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to get current task")) {
+            verifyActiveContainer();
+            final List<UiObject2> taskViews = mLauncher.getDevice().findObjects(
+                    mLauncher.getLauncherObjectSelector("snapshot"));
+            mLauncher.assertNotEquals("Unable to find a task", 0, taskViews.size());
 
-        // taskViews contains up to 3 task views: the 'main' (having the widest visible
-        // part) one in the center, and parts of its right and left siblings. Find the
-        // main task view by its width.
-        final UiObject2 widestTask = Collections.max(taskViews,
-                (t1, t2) -> Integer.compare(t1.getVisibleBounds().width(),
-                        t2.getVisibleBounds().width()));
+            // taskViews contains up to 3 task views: the 'main' (having the widest visible
+            // part) one in the center, and parts of its right and left siblings. Find the
+            // main task view by its width.
+            final UiObject2 widestTask = Collections.max(taskViews,
+                    (t1, t2) -> Integer.compare(t1.getVisibleBounds().width(),
+                            t2.getVisibleBounds().width()));
 
-        return new OverviewTask(mLauncher, widestTask, this);
+            return new OverviewTask(mLauncher, widestTask, this);
+        }
     }
 }
\ No newline at end of file
diff --git a/tests/tapl/com/android/launcher3/tapl/Home.java b/tests/tapl/com/android/launcher3/tapl/Home.java
index f8bd85a..20c116c 100644
--- a/tests/tapl/com/android/launcher3/tapl/Home.java
+++ b/tests/tapl/com/android/launcher3/tapl/Home.java
@@ -48,8 +48,14 @@
     @NonNull
     @Override
     public Overview switchToOverview() {
-        verifyActiveContainer();
-        goToOverviewUnchecked(OVERVIEW_STATE_ORDINAL);
-        return new Overview(mLauncher);
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to switch from home to overview")) {
+            verifyActiveContainer();
+            goToOverviewUnchecked(OVERVIEW_STATE_ORDINAL);
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
+                    "performed the switch action")) {
+                return new Overview(mLauncher);
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/tests/tapl/com/android/launcher3/tapl/Launchable.java b/tests/tapl/com/android/launcher3/tapl/Launchable.java
index 481281a..7a2b7af 100644
--- a/tests/tapl/com/android/launcher3/tapl/Launchable.java
+++ b/tests/tapl/com/android/launcher3/tapl/Launchable.java
@@ -59,10 +59,10 @@
     private Background launch(String errorMessage, BySelector selector) {
         LauncherInstrumentation.log("Launchable.launch before click " +
                 mObject.getVisibleCenter());
-        LauncherInstrumentation.assertTrue(
+        mLauncher.assertTrue(
                 "Launching an app didn't open a new window: " + mObject.getText(),
                 mObject.clickAndWait(Until.newWindow(), LauncherInstrumentation.WAIT_TIME_MS));
-        LauncherInstrumentation.assertTrue(
+        mLauncher.assertTrue(
                 "App didn't start: " + errorMessage,
                 mLauncher.getDevice().wait(Until.hasObject(selector),
                         LauncherInstrumentation.WAIT_TIME_MS));
@@ -79,6 +79,9 @@
                 this,
                 new Point(device.getDisplayWidth() / 2, device.getDisplayHeight() / 2),
                 DRAG_SPEED);
-        return new Workspace(mLauncher);
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "dragged launchable to workspace")) {
+            return new Workspace(mLauncher);
+        }
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index f7de2f9..19b368f 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -54,6 +54,8 @@
 
 import java.io.IOException;
 import java.lang.ref.WeakReference;
+import java.util.Deque;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.TimeoutException;
 
@@ -91,11 +93,16 @@
          * @return UI object for the container.
          */
         final UiObject2 verifyActiveContainer() {
-            assertTrue("Attempt to use a stale container", this == sActiveContainer.get());
+            mLauncher.assertTrue("Attempt to use a stale container",
+                    this == sActiveContainer.get());
             return mLauncher.verifyContainerType(getContainerType());
         }
     }
 
+    interface Closable extends AutoCloseable {
+        void close();
+    }
+
     private static final String WORKSPACE_RES_ID = "workspace";
     private static final String APPS_RES_ID = "apps_view";
     private static final String OVERVIEW_RES_ID = "overview_panel";
@@ -109,6 +116,7 @@
     private final Instrumentation mInstrumentation;
     private int mExpectedRotation = Surface.ROTATION_0;
     private final Uri mTestProviderUri;
+    private final Deque<String> mDiagnosticContext = new LinkedList<>();
 
     /**
      * Constructs the root of TAPL hierarchy. You get all other objects from it.
@@ -180,31 +188,38 @@
         Log.d(TAG, message);
     }
 
-    private static void fail(String message) {
-        Assert.fail("http://go/tapl : " + message);
+    Closable addContextLayer(String piece) {
+        mDiagnosticContext.addLast(piece);
+        return () -> mDiagnosticContext.removeLast();
     }
 
-    static void assertTrue(String message, boolean condition) {
+    private void fail(String message) {
+        final String ctxt = mDiagnosticContext.isEmpty() ? "" : String.join(", ",
+                mDiagnosticContext) + "; ";
+        Assert.fail("http://go/tapl : " + ctxt + message);
+    }
+
+    void assertTrue(String message, boolean condition) {
         if (!condition) {
             fail(message);
         }
     }
 
-    static void assertNotNull(String message, Object object) {
+    void assertNotNull(String message, Object object) {
         assertTrue(message, object != null);
     }
 
-    static private void failEquals(String message, Object actual) {
+    private void failEquals(String message, Object actual) {
         fail(message + ". " + "Actual: " + actual);
     }
 
-    static private void assertEquals(String message, int expected, int actual) {
+    private void assertEquals(String message, int expected, int actual) {
         if (expected != actual) {
             fail(message + " expected: " + expected + " but was: " + actual);
         }
     }
 
-    static void assertNotEquals(String message, int unexpected, int actual) {
+    void assertNotEquals(String message, int unexpected, int actual) {
         if (unexpected == actual) {
             failEquals(message, actual);
         }
@@ -222,49 +237,52 @@
                         (mDevice.findObject(By.res(SYSTEMUI_PACKAGE, "recent_apps")) == null));
         log("verifyContainerType: " + containerType);
 
-        switch (containerType) {
-            case WORKSPACE: {
-                waitForLauncherObject(APPS_RES_ID);
-                waitUntilGone(OVERVIEW_RES_ID);
-                waitUntilGone(WIDGETS_RES_ID);
-                return waitForLauncherObject(WORKSPACE_RES_ID);
-            }
-            case WIDGETS: {
-                waitUntilGone(WORKSPACE_RES_ID);
-                waitUntilGone(APPS_RES_ID);
-                waitUntilGone(OVERVIEW_RES_ID);
-                return waitForLauncherObject(WIDGETS_RES_ID);
-            }
-            case ALL_APPS: {
-                waitUntilGone(WORKSPACE_RES_ID);
-                waitUntilGone(OVERVIEW_RES_ID);
-                waitUntilGone(WIDGETS_RES_ID);
-                return waitForLauncherObject(APPS_RES_ID);
-            }
-            case OVERVIEW: {
-                if (mDevice.isNaturalOrientation()) {
+        try (Closable c = addContextLayer(
+                "but the current state is not " + containerType.name())) {
+            switch (containerType) {
+                case WORKSPACE: {
                     waitForLauncherObject(APPS_RES_ID);
-                } else {
-                    waitUntilGone(APPS_RES_ID);
+                    waitUntilGone(OVERVIEW_RES_ID);
+                    waitUntilGone(WIDGETS_RES_ID);
+                    return waitForLauncherObject(WORKSPACE_RES_ID);
                 }
-                // Fall through
-            }
-            case BASE_OVERVIEW: {
-                waitUntilGone(WORKSPACE_RES_ID);
-                waitUntilGone(WIDGETS_RES_ID);
+                case WIDGETS: {
+                    waitUntilGone(WORKSPACE_RES_ID);
+                    waitUntilGone(APPS_RES_ID);
+                    waitUntilGone(OVERVIEW_RES_ID);
+                    return waitForLauncherObject(WIDGETS_RES_ID);
+                }
+                case ALL_APPS: {
+                    waitUntilGone(WORKSPACE_RES_ID);
+                    waitUntilGone(OVERVIEW_RES_ID);
+                    waitUntilGone(WIDGETS_RES_ID);
+                    return waitForLauncherObject(APPS_RES_ID);
+                }
+                case OVERVIEW: {
+                    if (mDevice.isNaturalOrientation()) {
+                        waitForLauncherObject(APPS_RES_ID);
+                    } else {
+                        waitUntilGone(APPS_RES_ID);
+                    }
+                    // Fall through
+                }
+                case BASE_OVERVIEW: {
+                    waitUntilGone(WORKSPACE_RES_ID);
+                    waitUntilGone(WIDGETS_RES_ID);
 
-                return waitForLauncherObject(OVERVIEW_RES_ID);
+                    return waitForLauncherObject(OVERVIEW_RES_ID);
+                }
+                case BACKGROUND: {
+                    waitUntilGone(WORKSPACE_RES_ID);
+                    waitUntilGone(APPS_RES_ID);
+                    waitUntilGone(OVERVIEW_RES_ID);
+                    waitUntilGone(WIDGETS_RES_ID);
+                    return null;
+                }
+                default:
+                    fail("Invalid state: " + containerType);
+                    return null;
             }
-            case BACKGROUND: {
-                waitUntilGone(WORKSPACE_RES_ID);
-                waitUntilGone(APPS_RES_ID);
-                waitUntilGone(OVERVIEW_RES_ID);
-                waitUntilGone(WIDGETS_RES_ID);
-                return null;
-            }
-            default:
-                fail("Invalid state: " + containerType);
-                return null;
         }
     }
 
@@ -356,7 +374,9 @@
      */
     @NonNull
     public Workspace getWorkspace() {
-        return new Workspace(this);
+        try (LauncherInstrumentation.Closable c = addContextLayer("want to get workspace object")) {
+            return new Workspace(this);
+        }
     }
 
     /**
@@ -378,7 +398,9 @@
      */
     @NonNull
     public Widgets getAllWidgets() {
-        return new Widgets(this);
+        try (LauncherInstrumentation.Closable c = addContextLayer("want to get widgets")) {
+            return new Widgets(this);
+        }
     }
 
     /**
@@ -389,7 +411,9 @@
      */
     @NonNull
     public Overview getOverview() {
-        return new Overview(this);
+        try (LauncherInstrumentation.Closable c = addContextLayer("want to get overview")) {
+            return new Overview(this);
+        }
     }
 
     /**
@@ -413,7 +437,9 @@
      */
     @NonNull
     public AllApps getAllApps() {
-        return new AllApps(this);
+        try (LauncherInstrumentation.Closable c = addContextLayer("want to get all apps object")) {
+            return new AllApps(this);
+        }
     }
 
     /**
@@ -426,7 +452,9 @@
      */
     @NonNull
     public AllAppsFromOverview getAllAppsFromOverview() {
-        return new AllAppsFromOverview(this);
+        try (LauncherInstrumentation.Closable c = addContextLayer("want to get all apps object")) {
+            return new AllAppsFromOverview(this);
+        }
     }
 
     void waitUntilGone(String resId) {
diff --git a/tests/tapl/com/android/launcher3/tapl/Overview.java b/tests/tapl/com/android/launcher3/tapl/Overview.java
index 5c8d5eb..41932c8 100644
--- a/tests/tapl/com/android/launcher3/tapl/Overview.java
+++ b/tests/tapl/com/android/launcher3/tapl/Overview.java
@@ -20,11 +20,11 @@
 
 import android.graphics.Point;
 
-import com.android.launcher3.tapl.LauncherInstrumentation.ContainerType;
-
 import androidx.annotation.NonNull;
 import androidx.test.uiautomator.UiObject2;
 
+import com.android.launcher3.tapl.LauncherInstrumentation.ContainerType;
+
 /**
  * Overview pane.
  */
@@ -47,14 +47,20 @@
      */
     @NonNull
     public AllAppsFromOverview switchToAllApps() {
-        verifyActiveContainer();
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to switch from overview to all apps")) {
+            verifyActiveContainer();
 
-        // Swipe from navbar to the top.
-        final UiObject2 navBar = mLauncher.waitForSystemUiObject("navigation_bar_frame");
-        final Point start = navBar.getVisibleCenter();
-        LauncherInstrumentation.log("Overview.switchToAllApps before swipe");
-        mLauncher.swipe(start.x, start.y, start.x, 0, ALL_APPS_STATE_ORDINAL);
+            // Swipe from navbar to the top.
+            final UiObject2 navBar = mLauncher.waitForSystemUiObject("navigation_bar_frame");
+            final Point start = navBar.getVisibleCenter();
+            LauncherInstrumentation.log("Overview.switchToAllApps before swipe");
+            mLauncher.swipe(start.x, start.y, start.x, 0, ALL_APPS_STATE_ORDINAL);
 
-        return new AllAppsFromOverview(mLauncher);
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
+                    "swiped all way up from overview")) {
+                return new AllAppsFromOverview(mLauncher);
+            }
+        }
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
index 7ccd49b..b966851 100644
--- a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
+++ b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
@@ -44,10 +44,13 @@
      * Swipes the task up.
      */
     public void dismiss() {
-        verifyActiveContainer();
-        // Dismiss the task via flinging it up.
-        mTask.fling(Direction.DOWN, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        mLauncher.waitForIdle();
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to dismiss a task")) {
+            verifyActiveContainer();
+            // Dismiss the task via flinging it up.
+            mTask.fling(Direction.DOWN, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            mLauncher.waitForIdle();
+        }
     }
 
     /**
@@ -55,7 +58,7 @@
      */
     public Background open() {
         verifyActiveContainer();
-        LauncherInstrumentation.assertTrue("Launching task didn't open a new window: " +
+        mLauncher.assertTrue("Launching task didn't open a new window: " +
                         mTask.getParent().getContentDescription(),
                 mTask.clickAndWait(Until.newWindow(), LauncherInstrumentation.WAIT_TIME_MS));
         return new Background(mLauncher);
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 89affd1..b6938d6 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -34,21 +34,33 @@
      * Flings forward (down) and waits the fling's end.
      */
     public void flingForward() {
-        final UiObject2 widgetsContainer = verifyActiveContainer();
-        widgetsContainer.setGestureMargin(100);
-        widgetsContainer.fling(Direction.DOWN, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        verifyActiveContainer();
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to fling forward in widgets")) {
+            final UiObject2 widgetsContainer = verifyActiveContainer();
+            widgetsContainer.setGestureMargin(100);
+            widgetsContainer.fling(Direction.DOWN,
+                    (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer("flung forward")) {
+                verifyActiveContainer();
+            }
+        }
     }
 
     /**
      * Flings backward (up) and waits the fling's end.
      */
     public void flingBackward() {
-        final UiObject2 widgetsContainer = verifyActiveContainer();
-        widgetsContainer.setGestureMargin(100);
-        widgetsContainer.fling(Direction.UP, (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
-        mLauncher.waitForIdle();
-        verifyActiveContainer();
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to fling backwards in widgets")) {
+            final UiObject2 widgetsContainer = verifyActiveContainer();
+            widgetsContainer.setGestureMargin(100);
+            widgetsContainer.fling(Direction.UP,
+                    (int) (FLING_SPEED * mLauncher.getDisplayDensity()));
+            mLauncher.waitForIdle();
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer("flung back")) {
+                verifyActiveContainer();
+            }
+        }
     }
 
     @Override
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 6868e1e..9979f50 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -50,21 +50,27 @@
      */
     @NonNull
     public AllApps switchToAllApps() {
-        verifyActiveContainer();
-        final UiObject2 hotseat = mHotseat;
-        final Point start = hotseat.getVisibleCenter();
-        final int swipeHeight = mLauncher.getTestInfo(
-                TestProtocol.REQUEST_HOME_TO_ALL_APPS_SWIPE_HEIGHT).
-                getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
-        mLauncher.swipe(
-                start.x,
-                start.y,
-                start.x,
-                start.y - swipeHeight - mLauncher.getTouchSlop(),
-                ALL_APPS_STATE_ORDINAL
-        );
+        try(LauncherInstrumentation.Closable c =
+                    mLauncher.addContextLayer("want to switch from workspace to all apps")) {
+            verifyActiveContainer();
+            final UiObject2 hotseat = mHotseat;
+            final Point start = hotseat.getVisibleCenter();
+            final int swipeHeight = mLauncher.getTestInfo(
+                    TestProtocol.REQUEST_HOME_TO_ALL_APPS_SWIPE_HEIGHT).
+                    getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+            mLauncher.swipe(
+                    start.x,
+                    start.y,
+                    start.x,
+                    start.y - swipeHeight - mLauncher.getTouchSlop(),
+                    ALL_APPS_STATE_ORDINAL
+            );
 
-        return new AllApps(mLauncher);
+            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
+                    "swiped to all apps")) {
+                return new AllApps(mLauncher);
+            }
+        }
     }
 
     /**
@@ -75,9 +81,13 @@
      */
     @Nullable
     public AppIcon tryGetWorkspaceAppIcon(String appName) {
-        final UiObject2 workspace = verifyActiveContainer();
-        final UiObject2 icon = workspace.findObject(AppIcon.getAppIconSelector(appName, mLauncher));
-        return icon != null ? new AppIcon(mLauncher, icon) : null;
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to get a workspace icon")) {
+            final UiObject2 workspace = verifyActiveContainer();
+            final UiObject2 icon = workspace.findObject(
+                    AppIcon.getAppIconSelector(appName, mLauncher));
+            return icon != null ? new AppIcon(mLauncher, icon) : null;
+        }
     }
 
 
@@ -165,7 +175,9 @@
     public Widgets openAllWidgets() {
         verifyActiveContainer();
         mLauncher.getDevice().pressKeyCode(KeyEvent.KEYCODE_W, KeyEvent.META_CTRL_ON);
-        return new Widgets(mLauncher);
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer("pressed Ctrl+W")) {
+            return new Widgets(mLauncher);
+        }
     }
 
     @Override