1 files changed, 11 insertions, 5 deletions
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index dc72bda9a..5eb5f1992 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -27,18 +27,20 @@ import android.view.View;
 
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.UiObject2;
 
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.Workspace;
-import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.testcomponent.WidgetConfigActivity;
 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.Before;
 import org.junit.Ignore;
@@ -69,6 +71,7 @@ public class AddConfigWidgetTest extends AbstractLauncherUiTest {
     }
 
     @Test
+    // Convert test to TAPL b/131116002
     public void testWidgetConfig() throws Throwable {
         runTest(false, true);
     }
@@ -80,6 +83,7 @@ public class AddConfigWidgetTest extends AbstractLauncherUiTest {
     }
 
     @Test
+    // Convert test to TAPL b/131116002
     public void testConfigCancelled() throws Throwable {
         runTest(false, false);
     }
@@ -100,13 +104,15 @@ public class AddConfigWidgetTest extends AbstractLauncherUiTest {
         clearHomescreen();
         mActivityMonitor.startLauncher();
 
-        final Widgets widgets = mLauncher.getWorkspace().openAllWidgets();
+        // Open widget tray and wait for load complete.
+        final UiObject2 widgetContainer = TestViewHelpers.openWidgetsTray();
+        Wait.atMost(null, Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT);
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        widgets.
-                getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace();
+        UiObject2 widget = scrollAndFind(widgetContainer, By.clazz(WidgetCell.class)
+                .hasDescendant(By.text(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))));
+        TestViewHelpers.dragToWorkspace(widget, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();