diff --git a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
index 793af48..a3c70ec 100644
--- a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
@@ -45,6 +45,7 @@
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
 import com.android.launcher3.pm.InstallSessionHelper;
+import com.android.launcher3.tapl.Widget;
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
@@ -267,8 +268,10 @@
     }
 
     private void verifyWidgetPresent(LauncherAppWidgetProviderInfo info) {
+        final Widget widget = mLauncher.getWorkspace().tryGetWidget(info.label, DEFAULT_UI_TIMEOUT);
+        if (widget == null) mLauncher.dumpViewHierarchy(); // b/152645831
         assertTrue("Widget is not present",
-                mLauncher.getWorkspace().tryGetWidget(info.label, DEFAULT_UI_TIMEOUT) != null);
+                widget != null);
     }
 
     private void verifyPendingWidgetPresent() {
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index d894843..4648c2e 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -316,7 +316,7 @@
         };
     }
 
-    private void dumpViewHierarchy() {
+    public void dumpViewHierarchy() {
         final ByteArrayOutputStream stream = new ByteArrayOutputStream();
         try {
             mDevice.dumpWindowHierarchy(stream);
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 3f5dc8d..e730a80 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -281,16 +281,22 @@
 
     @Nullable
     public Widget tryGetWidget(String label, long timeout) {
-        final UiObject2 widget = mLauncher.tryWaitForLauncherObject(
-                By.clazz("com.android.launcher3.widget.LauncherAppWidgetHostView").desc(label),
-                timeout);
-        return widget != null ? new Widget(mLauncher, widget) : null;
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "getting widget " + label + " on workspace with timeout " + timeout)) {
+            final UiObject2 widget = mLauncher.tryWaitForLauncherObject(
+                    By.clazz("com.android.launcher3.widget.LauncherAppWidgetHostView").desc(label),
+                    timeout);
+            return widget != null ? new Widget(mLauncher, widget) : null;
+        }
     }
 
     @Nullable
     public Widget tryGetPendingWidget(long timeout) {
-        final UiObject2 widget = mLauncher.tryWaitForLauncherObject(
-                By.clazz("com.android.launcher3.widget.PendingAppWidgetHostView"), timeout);
-        return widget != null ? new Widget(mLauncher, widget) : null;
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "getting pending widget on workspace with timeout " + timeout)) {
+            final UiObject2 widget = mLauncher.tryWaitForLauncherObject(
+                    By.clazz("com.android.launcher3.widget.PendingAppWidgetHostView"), timeout);
+            return widget != null ? new Widget(mLauncher, widget) : null;
+        }
     }
 }
\ No newline at end of file