1 files changed, 23 insertions, 16 deletions
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index 3c92c578d..19f7db72b 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -15,42 +15,46 @@
  */
 package com.android.launcher3.ui.widget;
 
+import android.support.test.filters.LargeTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.support.test.uiautomator.By;
 import android.support.test.uiautomator.UiObject2;
-import android.test.suitebuilder.annotation.LargeTest;
 import android.view.View;
 
 import com.android.launcher3.ItemInfo;
-import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.Workspace.ItemOperator;
-import com.android.launcher3.ui.LauncherInstrumentationTestCase;
+import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Condition;
 import com.android.launcher3.util.Wait;
+import com.android.launcher3.util.rule.LauncherActivityRule;
+import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.WidgetCell;
 
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertTrue;
+
 /**
  * Test to add widget from widget tray
  */
 @LargeTest
-public class AddWidgetTest extends LauncherInstrumentationTestCase {
-
-    private LauncherAppWidgetProviderInfo widgetInfo;
+@RunWith(AndroidJUnit4.class)
+public class AddWidgetTest extends AbstractLauncherUiTest {
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        grantWidgetPermission();
-
-        widgetInfo = findWidgetProvider(false /* hasConfigureScreen */);
-    }
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grandWidgetBind();
 
+    @Test
     public void testDragIcon_portrait() throws Throwable {
         lockRotation(true);
         performTest();
     }
 
+    @Test
     public void testDragIcon_landscape() throws Throwable {
         lockRotation(false);
         performTest();
@@ -58,7 +62,10 @@ public class AddWidgetTest extends LauncherInstrumentationTestCase {
 
     private void performTest() throws Throwable {
         clearHomescreen();
-        Launcher launcher = startLauncher();
+        mActivityMonitor.startLauncher();
+
+        final LauncherAppWidgetProviderInfo widgetInfo =
+                findWidgetProvider(false /* hasConfigureScreen */);
 
         // Open widget tray and wait for load complete.
         final UiObject2 widgetContainer = openWidgetsTray();
@@ -69,12 +76,12 @@ public class AddWidgetTest extends LauncherInstrumentationTestCase {
                 .hasDescendant(By.text(widgetInfo.getLabel(mTargetContext.getPackageManager()))));
         dragToWorkspace(widget, false);
 
-        assertNotNull(launcher.getWorkspace().getFirstMatch(new ItemOperator() {
+        assertTrue(mActivityMonitor.itemExists(new ItemOperator() {
             @Override
             public boolean evaluate(ItemInfo info, View view) {
                 return info instanceof LauncherAppWidgetInfo &&
                         ((LauncherAppWidgetInfo) info).providerName.equals(widgetInfo.provider);
             }
-        }));
+        }).call());
     }
 }