1 files changed, 24 insertions, 8 deletions
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index 4529a80d1..0061568cc 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -19,12 +19,20 @@ import static org.junit.Assert.assertTrue;
 
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.UiObject2;
+import android.view.View;
 
+import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.Workspace.ItemOperator;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.Condition;
+import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ShellCommandRule;
+import com.android.launcher3.widget.WidgetCell;
 
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -53,6 +61,7 @@ public class AddWidgetTest extends AbstractLauncherUiTest {
         performTest();
     }
 
+    // Convert to TAPL b/131116002
     private void performTest() throws Throwable {
         clearHomescreen();
         mActivityMonitor.startLauncher();
@@ -60,15 +69,22 @@ public class AddWidgetTest extends AbstractLauncherUiTest {
         final LauncherAppWidgetProviderInfo widgetInfo =
                 TestViewHelpers.findWidgetProvider(this, false /* hasConfigureScreen */);
 
-        mLauncher.
-                getWorkspace().
-                openAllWidgets().
-                getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace();
+        // Open widget tray and wait for load complete.
+        final UiObject2 widgetContainer = TestViewHelpers.openWidgetsTray();
+        Wait.atMost(null, Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT);
 
-        assertTrue(mActivityMonitor.itemExists(
-                (info, view) -> info instanceof LauncherAppWidgetInfo &&
+        // Drag widget to homescreen
+        UiObject2 widget = scrollAndFind(widgetContainer, By.clazz(WidgetCell.class)
+                .hasDescendant(By.text(widgetInfo.getLabel(mTargetContext.getPackageManager()))));
+        TestViewHelpers.dragToWorkspace(widget, false);
+
+        assertTrue(mActivityMonitor.itemExists(new ItemOperator() {
+            @Override
+            public boolean evaluate(ItemInfo info, View view) {
+                return info instanceof LauncherAppWidgetInfo &&
                         ((LauncherAppWidgetInfo) info).providerName.getClassName().equals(
-                                widgetInfo.provider.getClassName())).call());
+                                widgetInfo.provider.getClassName());
+            }
+        }).call());
     }
 }