diff --git a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
index 21eb3fb..26f6ec3 100644
--- a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
+++ b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
@@ -71,10 +71,6 @@
         return result;
     }
 
-    public static boolean matches(AppInfo info, String query) {
-        return matches(info, query, StringMatcher.getInstance());
-    }
-
     public static boolean matches(AppInfo info, String query, StringMatcher matcher) {
         int queryLength = query.length();
 
diff --git a/src/com/android/launcher3/logging/FileLog.java b/src/com/android/launcher3/logging/FileLog.java
index 4c83e9a..f7f8ef1 100644
--- a/src/com/android/launcher3/logging/FileLog.java
+++ b/src/com/android/launcher3/logging/FileLog.java
@@ -29,6 +29,8 @@
  */
 public final class FileLog {
 
+    protected static final boolean ENABLED =
+            FeatureFlags.IS_DOGFOOD_BUILD || Utilities.IS_DEBUG_DEVICE;
     private static final String FILE_NAME_PREFIX = "log-";
     private static final DateFormat DATE_FORMAT =
             DateFormat.getDateTimeInstance(DateFormat.SHORT, DateFormat.SHORT);
@@ -39,7 +41,7 @@
     private static File sLogsDirectory = null;
 
     public static void setDir(File logsDir) {
-        if (FeatureFlags.IS_DOGFOOD_BUILD || Utilities.IS_DEBUG_DEVICE) {
+        if (ENABLED) {
             synchronized (DATE_FORMAT) {
                 // If the target directory changes, stop any active thread.
                 if (sHandler != null && !logsDir.equals(sLogsDirectory)) {
@@ -76,7 +78,7 @@
     }
 
     public static void print(String tag, String msg, Exception e) {
-        if (!FeatureFlags.IS_DOGFOOD_BUILD) {
+        if (!ENABLED) {
             return;
         }
         String out = String.format("%s %s %s", DATE_FORMAT.format(new Date()), tag, msg);
@@ -102,7 +104,7 @@
      * @param out if not null, all the persisted logs are copied to the writer.
      */
     public static void flushAll(PrintWriter out) throws InterruptedException {
-        if (!FeatureFlags.IS_DOGFOOD_BUILD) {
+        if (!ENABLED) {
             return;
         }
         CountDownLatch latch = new CountDownLatch(1);
@@ -135,7 +137,7 @@
 
         @Override
         public boolean handleMessage(Message msg) {
-            if (sLogsDirectory == null || !FeatureFlags.IS_DOGFOOD_BUILD) {
+            if (sLogsDirectory == null || !ENABLED) {
                 return true;
             }
             switch (msg.what) {
diff --git a/tests/src/com/android/launcher3/InvariantDeviceProfileTest.java b/tests/src/com/android/launcher3/InvariantDeviceProfileTest.java
deleted file mode 100644
index 230d623..0000000
--- a/tests/src/com/android/launcher3/InvariantDeviceProfileTest.java
+++ /dev/null
@@ -1,119 +0,0 @@
-/*
- * Copyright (C) 2015 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License"); you may not
- * use this file except in compliance with the License. You may obtain a copy of
- * the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
- * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
- * License for the specific language governing permissions and limitations under
- * the License.
- */
-package com.android.launcher3;
-
-import android.content.res.Resources;
-import android.graphics.Point;
-import android.graphics.PointF;
-import android.graphics.Rect;
-import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
-import android.util.Log;
-
-import java.util.ArrayList;
-
-/**
- * Tests the {@link DeviceProfile} and {@link InvariantDeviceProfile}.
- */
-@SmallTest
-public class InvariantDeviceProfileTest extends AndroidTestCase {
-
-    private static final String TAG = "DeviceProfileTest";
-    private static final boolean DEBUG = false;
-
-    private InvariantDeviceProfile mInvariantProfile;
-    private ArrayList<InvariantDeviceProfile> mPredefinedDeviceProfiles;
-
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        mInvariantProfile = new InvariantDeviceProfile(getContext());
-        mPredefinedDeviceProfiles = mInvariantProfile.getPredefinedDeviceProfiles(getContext());
-    }
-
-    @Override
-    protected void tearDown() throws Exception {
-        // Nothing to tear down as this class only tests static methods.
-    }
-
-    public void testFindClosestDeviceProfile2() {
-        for (InvariantDeviceProfile idf: mPredefinedDeviceProfiles) {
-            ArrayList<InvariantDeviceProfile> predefinedProfilesCopy =
-                    new ArrayList<>(mPredefinedDeviceProfiles);
-            ArrayList<InvariantDeviceProfile> closestProfiles =
-                    mInvariantProfile.findClosestDeviceProfiles(
-                            idf.minWidthDps, idf.minHeightDps, predefinedProfilesCopy
-                    );
-            assertTrue(closestProfiles.get(0).equals(idf));
-        }
-    }
-
-    /**
-     * Used to print out how the invDistWeightedInterpolate works between device profiles to
-     * tweak the two constants that control how the interpolation curve is shaped.
-     */
-    public void testInvInterpolation() {
-
-        InvariantDeviceProfile p1 = mPredefinedDeviceProfiles.get(7); // e.g., Large Phone
-        InvariantDeviceProfile p2 = mPredefinedDeviceProfiles.get(8); // e.g., Nexus 7
-
-        ArrayList<PointF> pts = createInterpolatedPoints(
-                new PointF(p1.minWidthDps, p1.minHeightDps),
-                new PointF(p2.minWidthDps, p2.minHeightDps),
-                20f);
-
-        for (int i = 0; i < pts.size(); i++) {
-            ArrayList<InvariantDeviceProfile> closestProfiles =
-                    mInvariantProfile.findClosestDeviceProfiles(
-                            pts.get(i).x, pts.get(i).y, mPredefinedDeviceProfiles);
-            InvariantDeviceProfile result =
-                    mInvariantProfile.invDistWeightedInterpolate(
-                            pts.get(i).x, pts.get(i).y, closestProfiles);
-            if (DEBUG) {
-                Log.d(TAG, String.format("width x height = (%f, %f)] iconSize = %f",
-                        pts.get(i).x, pts.get(i).y, result.iconSize));
-            }
-        }
-    }
-
-    private ArrayList<PointF> createInterpolatedPoints(PointF a, PointF b, float numPts) {
-        ArrayList<PointF> result = new ArrayList<PointF>();
-        result.add(a);
-        for (float i = 1; i < numPts; i = i + 1.0f) {
-            result.add(new PointF((b.x * i +  a.x * (numPts - i)) / numPts,
-                    (b.y * i + a.y * (numPts - i)) / numPts));
-        }
-        result.add(b);
-        return result;
-    }
-
-    /**
-     * Ensures that system calls (e.g., WindowManager, DisplayMetrics) that require contexts are
-     * properly working to generate minimum width and height of the display.
-     */
-    public void test_hammerhead() {
-        if (!android.os.Build.DEVICE.equals("hammerhead")) {
-            return;
-        }
-        assertEquals(4, mInvariantProfile.numRows);
-        assertEquals(4, mInvariantProfile.numColumns);
-        assertEquals(5, mInvariantProfile.numHotseatIcons);
-    }
-
-    // Add more tests for other devices, however, running them once on a single device is enough
-    // for verifying that for a platform version, the WindowManager and DisplayMetrics is
-    // working as intended.
-}
diff --git a/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java b/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java
index 26ec69b..846a163 100644
--- a/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java
+++ b/tests/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithmTest.java
@@ -16,85 +16,78 @@
 package com.android.launcher3.allapps.search;
 
 import android.content.ComponentName;
-import android.test.InstrumentationTestCase;
+import android.support.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.AppInfo;
 import com.android.launcher3.Utilities;
 
-import java.util.ArrayList;
-import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Unit tests for {@link DefaultAppSearchAlgorithm}
  */
-public class DefaultAppSearchAlgorithmTest extends InstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class DefaultAppSearchAlgorithmTest {
+    private static final DefaultAppSearchAlgorithm.StringMatcher MATCHER =
+            DefaultAppSearchAlgorithm.StringMatcher.getInstance();
 
-    private List<AppInfo> mAppsList;
-    private DefaultAppSearchAlgorithm mAlgorithm;
-
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        mAppsList = new ArrayList<>();
-        getInstrumentation().runOnMainSync(new Runnable() {
-            @Override
-            public void run() {
-                mAlgorithm = new DefaultAppSearchAlgorithm(mAppsList);
-            }
-        });
-    }
-
+    @Test
     public void testMatches() {
-        assertTrue(mAlgorithm.matches(getInfo("white cow"), "cow"));
-        assertTrue(mAlgorithm.matches(getInfo("whiteCow"), "cow"));
-        assertTrue(mAlgorithm.matches(getInfo("whiteCOW"), "cow"));
-        assertTrue(mAlgorithm.matches(getInfo("whitecowCOW"), "cow"));
-        assertTrue(mAlgorithm.matches(getInfo("white2cow"), "cow"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("white cow"), "cow", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whiteCow"), "cow", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whiteCOW"), "cow", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whitecowCOW"), "cow", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("white2cow"), "cow", MATCHER));
 
-        assertFalse(mAlgorithm.matches(getInfo("whitecow"), "cow"));
-        assertFalse(mAlgorithm.matches(getInfo("whitEcow"), "cow"));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whitecow"), "cow", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whitEcow"), "cow", MATCHER));
 
-        assertTrue(mAlgorithm.matches(getInfo("whitecowCow"), "cow"));
-        assertTrue(mAlgorithm.matches(getInfo("whitecow cow"), "cow"));
-        assertFalse(mAlgorithm.matches(getInfo("whitecowcow"), "cow"));
-        assertFalse(mAlgorithm.matches(getInfo("whit ecowcow"), "cow"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whitecowCow"), "cow", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("whitecow cow"), "cow", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whitecowcow"), "cow", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("whit ecowcow"), "cow", MATCHER));
 
-        assertTrue(mAlgorithm.matches(getInfo("cats&dogs"), "dog"));
-        assertTrue(mAlgorithm.matches(getInfo("cats&Dogs"), "dog"));
-        assertTrue(mAlgorithm.matches(getInfo("cats&Dogs"), "&"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("cats&dogs"), "dog", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("cats&Dogs"), "dog", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("cats&Dogs"), "&", MATCHER));
 
-        assertTrue(mAlgorithm.matches(getInfo("2+43"), "43"));
-        assertFalse(mAlgorithm.matches(getInfo("2+43"), "3"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("2+43"), "43", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("2+43"), "3", MATCHER));
 
-        assertTrue(mAlgorithm.matches(getInfo("Q"), "q"));
-        assertTrue(mAlgorithm.matches(getInfo("  Q"), "q"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("Q"), "q", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("  Q"), "q", MATCHER));
 
         // match lower case words
-        assertTrue(mAlgorithm.matches(getInfo("elephant"), "e"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("elephant"), "e", MATCHER));
 
-        assertTrue(mAlgorithm.matches(getInfo("电子邮件"), "电"));
-        assertTrue(mAlgorithm.matches(getInfo("电子邮件"), "电子"));
-        assertFalse(mAlgorithm.matches(getInfo("电子邮件"), "子"));
-        assertFalse(mAlgorithm.matches(getInfo("电子邮件"), "邮件"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "电", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "电子", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "子", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("电子邮件"), "邮件", MATCHER));
 
-        assertFalse(mAlgorithm.matches(getInfo("Bot"), "ba"));
-        assertFalse(mAlgorithm.matches(getInfo("bot"), "ba"));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("Bot"), "ba", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("bot"), "ba", MATCHER));
     }
 
+    @Test
     public void testMatchesVN() {
         if (!Utilities.ATLEAST_NOUGAT) {
             return;
         }
-        assertTrue(mAlgorithm.matches(getInfo("다운로드"), "다"));
-        assertTrue(mAlgorithm.matches(getInfo("드라이브"), "드"));
-        assertTrue(mAlgorithm.matches(getInfo("다운로드 드라이브"), "ㄷ"));
-        assertTrue(mAlgorithm.matches(getInfo("운로 드라이브"), "ㄷ"));
-        assertTrue(mAlgorithm.matches(getInfo("abc"), "åbç"));
-        assertTrue(mAlgorithm.matches(getInfo("Alpha"), "ål"));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("다운로드"), "다", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("드라이브"), "드", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("다운로드 드라이브"), "ㄷ", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("운로 드라이브"), "ㄷ", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("abc"), "åbç", MATCHER));
+        assertTrue(DefaultAppSearchAlgorithm.matches(getInfo("Alpha"), "ål", MATCHER));
 
-        assertFalse(mAlgorithm.matches(getInfo("다운로드 드라이브"), "ㄷㄷ"));
-        assertFalse(mAlgorithm.matches(getInfo("로드라이브"), "ㄷ"));
-        assertFalse(mAlgorithm.matches(getInfo("abc"), "åç"));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("다운로드 드라이브"), "ㄷㄷ", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("로드라이브"), "ㄷ", MATCHER));
+        assertFalse(DefaultAppSearchAlgorithm.matches(getInfo("abc"), "åç", MATCHER));
     }
 
     private AppInfo getInfo(String title) {
diff --git a/tests/src/com/android/launcher3/logging/FileLogTest.java b/tests/src/com/android/launcher3/logging/FileLogTest.java
index c24cc3f..7048c28 100644
--- a/tests/src/com/android/launcher3/logging/FileLogTest.java
+++ b/tests/src/com/android/launcher3/logging/FileLogTest.java
@@ -37,6 +37,9 @@
     }
 
     public void testPrintLog() throws Exception {
+        if (!FileLog.ENABLED) {
+            return;
+        }
         FileLog.print("Testing", "hoolalala");
         StringWriter writer = new StringWriter();
         FileLog.flushAll(new PrintWriter(writer));
@@ -54,6 +57,9 @@
     }
 
     public void testOldFileTruncated() throws Exception {
+        if (!FileLog.ENABLED) {
+            return;
+        }
         FileLog.print("Testing", "hoolalala");
         StringWriter writer = new StringWriter();
         FileLog.flushAll(new PrintWriter(writer));
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
index 4c80902..ae15f08 100644
--- a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -20,9 +20,9 @@
 import org.mockito.ArgumentCaptor;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 
+import static org.mockito.Matchers.isNull;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -131,7 +131,7 @@
         // only info2 should be added because info was already added to the workspace
         // in setupWorkspaceWithHoles()
         verify(callbacks).bindAppsAdded(any(ArrayList.class), notAnimated.capture(),
-                animated.capture(), any(ArrayList.class));
+                animated.capture(), isNull(ArrayList.class));
         assertTrue(notAnimated.getValue().isEmpty());
 
         assertEquals(1, animated.getValue().size());
diff --git a/tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
similarity index 80%
rename from tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java
rename to tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 47b43f5..7fb5d85 100644
--- a/tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -15,29 +15,26 @@
  */
 package com.android.launcher3.ui;
 
+import android.app.Instrumentation;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
-import android.content.pm.ActivityInfo;
-import android.content.pm.PackageManager;
-import android.content.pm.ResolveInfo;
+import android.content.pm.LauncherActivityInfo;
 import android.graphics.Point;
-import android.os.ParcelFileDescriptor;
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.support.test.InstrumentationRegistry;
 import android.support.test.uiautomator.By;
 import android.support.test.uiautomator.BySelector;
 import android.support.test.uiautomator.Direction;
 import android.support.test.uiautomator.UiDevice;
 import android.support.test.uiautomator.UiObject2;
 import android.support.test.uiautomator.Until;
-import android.test.InstrumentationTestCase;
 import android.view.MotionEvent;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.LauncherSettings;
@@ -45,25 +42,26 @@
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
+import com.android.launcher3.compat.LauncherAppsCompat;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.testcomponent.AppWidgetNoConfig;
 import com.android.launcher3.testcomponent.AppWidgetWithConfig;
 import com.android.launcher3.util.ManagedProfileHeuristic;
 
-import java.io.BufferedReader;
-import java.io.FileInputStream;
-import java.io.IOException;
-import java.io.InputStreamReader;
-import java.util.ArrayList;
+import org.junit.Before;
+
 import java.util.Locale;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
 /**
  * Base class for all instrumentation tests providing various utility methods.
  */
-public class LauncherInstrumentationTestCase extends InstrumentationTestCase {
+public abstract class AbstractLauncherUiTest {
 
     public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 5;
@@ -71,16 +69,15 @@
     public static final long DEFAULT_UI_TIMEOUT = 3000;
     public static final long DEFAULT_WORKER_TIMEOUT_SECS = 5;
 
+    protected MainThreadExecutor mMainThreadExecutor = new MainThreadExecutor();
     protected UiDevice mDevice;
     protected Context mTargetContext;
     protected String mTargetPackage;
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
+    @Before
+    public void setUp() throws Exception {
         mDevice = UiDevice.getInstance(getInstrumentation());
-        mTargetContext = getInstrumentation().getTargetContext();
+        mTargetContext = InstrumentationRegistry.getTargetContext();
         mTargetPackage = mTargetContext.getPackageName();
     }
 
@@ -97,56 +94,15 @@
         }
     }
 
-    /**
-     * Starts the launcher activity in the target package and returns the Launcher instance.
-     */
-    protected Launcher startLauncher() {
-        return (Launcher) getInstrumentation().startActivitySync(getHomeIntent());
-    }
-
-    protected Intent getHomeIntent() {
-        return new Intent(Intent.ACTION_MAIN)
-                .addCategory(Intent.CATEGORY_HOME)
-                .setPackage(mTargetPackage)
-                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
-    }
-
-    /**
-     * Grants the launcher permission to bind widgets.
-     */
-    protected void grantWidgetPermission() throws IOException {
-        // Check bind widget permission
-        if (mTargetContext.getPackageManager().checkPermission(
-                mTargetPackage, android.Manifest.permission.BIND_APPWIDGET)
-                != PackageManager.PERMISSION_GRANTED) {
-            runShellCommand("appwidget grantbind --package " + mTargetPackage);
-        }
-    }
-
-    /**
-     * Sets the target launcher as default launcher.
-     */
-    protected void setDefaultLauncher() throws IOException {
-        ActivityInfo launcher = mTargetContext.getPackageManager()
-                .queryIntentActivities(getHomeIntent(), 0).get(0).activityInfo;
-        runShellCommand("cmd package set-home-activity " +
-                new ComponentName(launcher.packageName, launcher.name).flattenToString());
-    }
-
-    protected void runShellCommand(String command) throws IOException {
-        ParcelFileDescriptor pfd = getInstrumentation().getUiAutomation()
-                .executeShellCommand(command);
-
-        // Read the input stream fully.
-        FileInputStream fis = new ParcelFileDescriptor.AutoCloseInputStream(pfd);
-        while (fis.read() != -1);
-        fis.close();
+    protected Instrumentation getInstrumentation() {
+        return InstrumentationRegistry.getInstrumentation();
     }
 
     /**
      * Opens all apps and returns the recycler view
      */
     protected UiObject2 openAllApps() {
+        mDevice.waitForIdle();
         if (FeatureFlags.LAUNCHER3_ALL_APPS_PULL_UP) {
             // clicking on the page indicator brings up all apps tray on non tablets.
             findViewById(R.id.page_indicator).click();
@@ -262,7 +218,7 @@
 
     protected void resetLoaderState() {
         try {
-            runTestOnUiThread(new Runnable() {
+            mMainThreadExecutor.execute(new Runnable() {
                 @Override
                 public void run() {
                     ManagedProfileHeuristic.markExistingUsersForNoFolderCreation(mTargetContext);
@@ -279,7 +235,7 @@
      */
     protected <T> T getOnUiThread(final Callable<T> callback) {
         try {
-            return new MainThreadExecutor().submit(callback).get();
+            return mMainThreadExecutor.submit(callback).get();
         } catch (Exception e) {
             throw new RuntimeException(e);
         }
@@ -315,6 +271,10 @@
         return By.res(mTargetPackage, name);
     }
 
+    protected LauncherActivityInfo getSettingsApp() {
+        return LauncherAppsCompat.getInstance(mTargetContext)
+                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
+    }
 
     /**
      * Broadcast receiver which blocks until the result is received.
diff --git a/tests/src/com/android/launcher3/ui/AllAppsAppLaunchTest.java b/tests/src/com/android/launcher3/ui/AllAppsAppLaunchTest.java
index 0ced7cf..46343a3 100644
--- a/tests/src/com/android/launcher3/ui/AllAppsAppLaunchTest.java
+++ b/tests/src/com/android/launcher3/ui/AllAppsAppLaunchTest.java
@@ -1,52 +1,55 @@
 package com.android.launcher3.ui;
 
 import android.content.pm.LauncherActivityInfo;
-import android.os.Process;
+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.support.test.uiautomator.Until;
-import android.test.suitebuilder.annotation.LargeTest;
 
-import com.android.launcher3.compat.LauncherAppsCompat;
 import com.android.launcher3.util.Condition;
 import com.android.launcher3.util.Wait;
+import com.android.launcher3.util.rule.LauncherActivityRule;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertTrue;
 
 /**
  * Test for verifying apps is launched from all-apps
  */
 @LargeTest
-public class AllAppsAppLaunchTest extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class AllAppsAppLaunchTest extends AbstractLauncherUiTest {
 
-    private LauncherActivityInfo mSettingsApp;
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
-        mSettingsApp = LauncherAppsCompat.getInstance(mTargetContext)
-                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
-    }
-
+    @Test
     public void testAppLauncher_portrait() throws Exception {
         lockRotation(true);
         performTest();
     }
 
+    @Test
     public void testAppLauncher_landscape() throws Exception {
         lockRotation(false);
         performTest();
     }
 
     private void performTest() throws Exception {
-        startLauncher();
+        mActivityMonitor.startLauncher();
+
+        LauncherActivityInfo settingsApp = getSettingsApp();
 
         // Open all apps and wait for load complete
         final UiObject2 appsContainer = openAllApps();
         assertTrue(Wait.atMost(Condition.minChildCount(appsContainer, 2), DEFAULT_UI_TIMEOUT));
 
         // Open settings app and verify app launched
-        scrollAndFind(appsContainer, By.text(mSettingsApp.getLabel().toString())).click();
+        scrollAndFind(appsContainer, By.text(settingsApp.getLabel().toString())).click();
         assertTrue(mDevice.wait(Until.hasObject(By.pkg(
-                mSettingsApp.getComponentName().getPackageName()).depth(0)), DEFAULT_UI_TIMEOUT));
+                settingsApp.getComponentName().getPackageName()).depth(0)), DEFAULT_UI_TIMEOUT));
     }
 }
diff --git a/tests/src/com/android/launcher3/ui/AllAppsIconToHomeTest.java b/tests/src/com/android/launcher3/ui/AllAppsIconToHomeTest.java
index 9361750..00f30ad 100644
--- a/tests/src/com/android/launcher3/ui/AllAppsIconToHomeTest.java
+++ b/tests/src/com/android/launcher3/ui/AllAppsIconToHomeTest.java
@@ -1,58 +1,62 @@
 package com.android.launcher3.ui;
 
 import android.content.pm.LauncherActivityInfo;
-import android.os.Process;
+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.support.test.uiautomator.Until;
-import android.test.suitebuilder.annotation.LargeTest;
 
-import com.android.launcher3.compat.LauncherAppsCompat;
 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 org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertTrue;
 
 /**
  * Test for dragging an icon from all-apps to homescreen.
  */
 @LargeTest
-public class AllAppsIconToHomeTest extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class AllAppsIconToHomeTest extends AbstractLauncherUiTest {
 
-    private LauncherActivityInfo mSettingsApp;
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mDefaultLauncherRule = ShellCommandRule.setDefaultLauncher();
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        setDefaultLauncher();
-
-        mSettingsApp = LauncherAppsCompat.getInstance(mTargetContext)
-                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
-    }
-
+    @Test
     public void testDragIcon_portrait() throws Throwable {
         lockRotation(true);
         performTest();
     }
 
+    @Test
     public void testDragIcon_landscape() throws Throwable {
         lockRotation(false);
         performTest();
     }
 
     private void performTest() throws Throwable {
+        LauncherActivityInfo settingsApp = getSettingsApp();
+
         clearHomescreen();
-        startLauncher();
+        mActivityMonitor.startLauncher();
 
         // Open all apps and wait for load complete.
         final UiObject2 appsContainer = openAllApps();
         assertTrue(Wait.atMost(Condition.minChildCount(appsContainer, 2), DEFAULT_UI_TIMEOUT));
 
         // Drag icon to homescreen.
-        UiObject2 icon = scrollAndFind(appsContainer, By.text(mSettingsApp.getLabel().toString()));
+        UiObject2 icon = scrollAndFind(appsContainer, By.text(settingsApp.getLabel().toString()));
         dragToWorkspace(icon, true);
 
         // Verify that the icon works on homescreen.
-        mDevice.findObject(By.text(mSettingsApp.getLabel().toString())).click();
+        mDevice.findObject(By.text(settingsApp.getLabel().toString())).click();
         assertTrue(mDevice.wait(Until.hasObject(By.pkg(
-                mSettingsApp.getComponentName().getPackageName()).depth(0)), DEFAULT_UI_TIMEOUT));
+                settingsApp.getComponentName().getPackageName()).depth(0)), DEFAULT_UI_TIMEOUT));
     }
 }
diff --git a/tests/src/com/android/launcher3/ui/RotationPreferenceTest.java b/tests/src/com/android/launcher3/ui/RotationPreferenceTest.java
deleted file mode 100644
index e84ad04..0000000
--- a/tests/src/com/android/launcher3/ui/RotationPreferenceTest.java
+++ /dev/null
@@ -1,77 +0,0 @@
-package com.android.launcher3.ui;
-
-import android.content.SharedPreferences;
-import android.graphics.Rect;
-import android.support.test.uiautomator.UiDevice;
-import android.support.test.uiautomator.UiObject;
-import android.support.test.uiautomator.UiSelector;
-import android.test.suitebuilder.annotation.MediumTest;
-
-import com.android.launcher3.R;
-import com.android.launcher3.Utilities;
-
-/**
- * Test for auto rotate preference.
- */
-@MediumTest
-public class RotationPreferenceTest extends LauncherInstrumentationTestCase {
-
-    private SharedPreferences mPrefs;
-    private boolean mOriginalRotationValue;
-
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
-        mDevice = UiDevice.getInstance(getInstrumentation());
-        mTargetContext = getInstrumentation().getTargetContext();
-        mTargetPackage = mTargetContext.getPackageName();
-        mPrefs = Utilities.getPrefs(mTargetContext);
-        mOriginalRotationValue = mPrefs.getBoolean(Utilities.ALLOW_ROTATION_PREFERENCE_KEY, false);
-    }
-
-    @Override
-    protected void tearDown() throws Exception {
-        setRotationEnabled(mOriginalRotationValue);
-        super.tearDown();
-    }
-
-    public void testRotation_disabled() throws Exception {
-        if (mTargetContext.getResources().getBoolean(R.bool.allow_rotation)) {
-            // This is a tablet. The test is only valid to mobile devices.
-            return;
-        }
-
-        setRotationEnabled(false);
-        mDevice.setOrientationRight();
-        startLauncher();
-
-        Rect hotseat = getHotseatBounds();
-        assertTrue(hotseat.width() > hotseat.height());
-    }
-
-    public void testRotation_enabled() throws Exception {
-        if (mTargetContext.getResources().getBoolean(R.bool.allow_rotation)) {
-            // This is a tablet. The test is only valid to mobile devices.
-            return;
-        }
-
-        setRotationEnabled(true);
-        mDevice.setOrientationRight();
-        startLauncher();
-
-        Rect hotseat = getHotseatBounds();
-        assertTrue(hotseat.width() < hotseat.height());
-    }
-
-    private void setRotationEnabled(boolean enabled) {
-        mPrefs.edit().putBoolean(Utilities.ALLOW_ROTATION_PREFERENCE_KEY, enabled).commit();
-    }
-
-    private Rect getHotseatBounds() throws Exception {
-        UiObject hotseat = mDevice.findObject(
-                new UiSelector().resourceId(mTargetPackage + ":id/hotseat"));
-        hotseat.waitForExists(6000);
-        return hotseat.getVisibleBounds();
-    }
-}
diff --git a/tests/src/com/android/launcher3/ui/ShortcutsLaunchTest.java b/tests/src/com/android/launcher3/ui/ShortcutsLaunchTest.java
index 3a0b613..a40ad7f 100644
--- a/tests/src/com/android/launcher3/ui/ShortcutsLaunchTest.java
+++ b/tests/src/com/android/launcher3/ui/ShortcutsLaunchTest.java
@@ -2,54 +2,58 @@
 
 import android.content.pm.LauncherActivityInfo;
 import android.graphics.Point;
-import android.os.Process;
+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.support.test.uiautomator.Until;
-import android.test.suitebuilder.annotation.LargeTest;
 import android.view.MotionEvent;
 
 import com.android.launcher3.R;
-import com.android.launcher3.compat.LauncherAppsCompat;
 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 org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Test for verifying that shortcuts are shown and can be launched after long pressing an app
  */
 @LargeTest
-public class ShortcutsLaunchTest extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class ShortcutsLaunchTest extends AbstractLauncherUiTest {
 
-    private LauncherActivityInfo mSettingsApp;
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mDefaultLauncherRule = ShellCommandRule.setDefaultLauncher();
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        setDefaultLauncher();
-
-        mSettingsApp = LauncherAppsCompat.getInstance(mTargetContext)
-                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
-    }
-
+    @Test
     public void testAppLauncher_portrait() throws Exception {
         lockRotation(true);
         performTest();
     }
 
+    @Test
     public void testAppLauncher_landscape() throws Exception {
         lockRotation(false);
         performTest();
     }
 
     private void performTest() throws Exception {
-        startLauncher();
+        mActivityMonitor.startLauncher();
+        LauncherActivityInfo settingsApp = getSettingsApp();
 
         // Open all apps and wait for load complete
         final UiObject2 appsContainer = openAllApps();
         assertTrue(Wait.atMost(Condition.minChildCount(appsContainer, 2), DEFAULT_UI_TIMEOUT));
 
         // Find settings app and verify shortcuts appear when long pressed
-        UiObject2 icon = scrollAndFind(appsContainer, By.text(mSettingsApp.getLabel().toString()));
+        UiObject2 icon = scrollAndFind(appsContainer, By.text(settingsApp.getLabel().toString()));
         // Press icon center until shortcuts appear
         Point iconCenter = icon.getVisibleCenter();
         sendPointer(MotionEvent.ACTION_DOWN, iconCenter);
@@ -63,7 +67,7 @@
                 .findObject(getSelectorForId(R.id.bubble_text));
         shortcut.click();
         assertTrue(mDevice.wait(Until.hasObject(By.pkg(
-                mSettingsApp.getComponentName().getPackageName())
+                settingsApp.getComponentName().getPackageName())
                 .text(shortcut.getText())), DEFAULT_UI_TIMEOUT));
     }
 }
diff --git a/tests/src/com/android/launcher3/ui/ShortcutsToHomeTest.java b/tests/src/com/android/launcher3/ui/ShortcutsToHomeTest.java
index 5d86d1e..434311d 100644
--- a/tests/src/com/android/launcher3/ui/ShortcutsToHomeTest.java
+++ b/tests/src/com/android/launcher3/ui/ShortcutsToHomeTest.java
@@ -2,40 +2,43 @@
 
 import android.content.pm.LauncherActivityInfo;
 import android.graphics.Point;
-import android.os.Process;
+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.support.test.uiautomator.Until;
-import android.test.suitebuilder.annotation.LargeTest;
 import android.view.MotionEvent;
 
 import com.android.launcher3.R;
-import com.android.launcher3.compat.LauncherAppsCompat;
 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 org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Test for dragging a deep shortcut to the home screen.
  */
 @LargeTest
-public class ShortcutsToHomeTest extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class ShortcutsToHomeTest extends AbstractLauncherUiTest {
 
-    private LauncherActivityInfo mSettingsApp;
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mDefaultLauncherRule = ShellCommandRule.setDefaultLauncher();
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        setDefaultLauncher();
-
-        mSettingsApp = LauncherAppsCompat.getInstance(mTargetContext)
-                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
-    }
-
+    @Test
     public void testDragIcon_portrait() throws Throwable {
         lockRotation(true);
         performTest();
     }
 
+    @Test
     public void testDragIcon_landscape() throws Throwable {
         lockRotation(false);
         performTest();
@@ -43,14 +46,16 @@
 
     private void performTest() throws Throwable {
         clearHomescreen();
-        startLauncher();
+        mActivityMonitor.startLauncher();
+
+        LauncherActivityInfo settingsApp  = getSettingsApp();
 
         // Open all apps and wait for load complete.
         final UiObject2 appsContainer = openAllApps();
         assertTrue(Wait.atMost(Condition.minChildCount(appsContainer, 2), DEFAULT_UI_TIMEOUT));
 
         // Find the app and long press it to show shortcuts.
-        UiObject2 icon = scrollAndFind(appsContainer, By.text(mSettingsApp.getLabel().toString()));
+        UiObject2 icon = scrollAndFind(appsContainer, By.text(settingsApp.getLabel().toString()));
         // Press icon center until shortcuts appear
         Point iconCenter = icon.getVisibleCenter();
         sendPointer(MotionEvent.ACTION_DOWN, iconCenter);
@@ -69,7 +74,7 @@
         // (the app opens and has the same text as the shortcut).
         mDevice.findObject(By.text(shortcutName)).click();
         assertTrue(mDevice.wait(Until.hasObject(By.pkg(
-                mSettingsApp.getComponentName().getPackageName())
+                settingsApp.getComponentName().getPackageName())
                 .text(shortcutName)), DEFAULT_UI_TIMEOUT));
     }
 }
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 0b4e34f..a5c2e69 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -15,75 +15,75 @@
  */
 package com.android.launcher3.ui.widget;
 
-import android.app.Activity;
-import android.app.Application;
 import android.appwidget.AppWidgetManager;
 import android.content.Intent;
+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.MainThreadExecutor;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.testcomponent.WidgetConfigActivity;
-import com.android.launcher3.ui.LauncherInstrumentationTestCase;
+import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Condition;
-import com.android.launcher3.util.SimpleActivityMonitor;
 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 java.util.concurrent.Callable;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Test to verify widget configuration is properly shown.
  */
 @LargeTest
-public class AddConfigWidgetTest extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class AddConfigWidgetTest extends AbstractLauncherUiTest {
+
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grandWidgetBind();
 
     private LauncherAppWidgetProviderInfo mWidgetInfo;
-    private SimpleActivityMonitor mActivityMonitor;
-    private MainThreadExecutor mMainThreadExecutor;
     private AppWidgetManager mAppWidgetManager;
 
     private int mWidgetId;
 
     @Override
-    protected void setUp() throws Exception {
+    @Before
+    public void setUp() throws Exception {
         super.setUp();
         mWidgetInfo = findWidgetProvider(true /* hasConfigureScreen */);
-        mActivityMonitor = new SimpleActivityMonitor();
-        ((Application) getInstrumentation().getTargetContext().getApplicationContext())
-                .registerActivityLifecycleCallbacks(mActivityMonitor);
-        mMainThreadExecutor = new MainThreadExecutor();
         mAppWidgetManager = AppWidgetManager.getInstance(mTargetContext);
-
-        grantWidgetPermission();
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        ((Application) getInstrumentation().getTargetContext().getApplicationContext())
-                .unregisterActivityLifecycleCallbacks(mActivityMonitor);
-        super.tearDown();
-    }
-
+    @Test
     public void testWidgetConfig() throws Throwable {
         runTest(false, true);
     }
 
+    @Test
     public void testWidgetConfig_rotate() throws Throwable {
         runTest(true, true);
     }
 
+    @Test
     public void testConfigCancelled() throws Throwable {
         runTest(false, false);
     }
 
+    @Test
     public void testConfigCancelled_rotate() throws Throwable {
         runTest(true, false);
     }
@@ -96,7 +96,7 @@
         lockRotation(true);
 
         clearHomescreen();
-        startLauncher();
+        mActivityMonitor.startLauncher();
 
         // Open widget tray and wait for load complete.
         final UiObject2 widgetContainer = openWidgetsTray();
@@ -146,11 +146,11 @@
      * Condition for searching widget id
      */
     private class WidgetSearchCondition extends Condition
-            implements Callable<Boolean>, Workspace.ItemOperator {
+            implements Workspace.ItemOperator {
 
         @Override
         public boolean isTrue() throws Throwable {
-            return mMainThreadExecutor.submit(this).get();
+            return mMainThreadExecutor.submit(mActivityMonitor.itemExists(this)).get();
         }
 
         @Override
@@ -159,21 +159,6 @@
                     ((LauncherAppWidgetInfo) info).providerName.equals(mWidgetInfo.provider) &&
                     ((LauncherAppWidgetInfo) info).appWidgetId == mWidgetId;
         }
-
-        @Override
-        public Boolean call() throws Exception {
-            // Find the resumed launcher
-            Launcher launcher = null;
-            for (Activity a : mActivityMonitor.resumed) {
-                if (a instanceof Launcher) {
-                    launcher = (Launcher) a;
-                }
-            }
-            if (launcher == null) {
-                return false;
-            }
-            return launcher.getWorkspace().getFirstMatch(this) != null;
-        }
     }
 
     /**
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index 3c92c57..19f7db7 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 {
+@RunWith(AndroidJUnit4.class)
+public class AddWidgetTest extends AbstractLauncherUiTest {
 
-    private LauncherAppWidgetProviderInfo widgetInfo;
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grandWidgetBind();
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        grantWidgetPermission();
-
-        widgetInfo = findWidgetProvider(false /* hasConfigureScreen */);
-    }
-
+    @Test
     public void testDragIcon_portrait() throws Throwable {
         lockRotation(true);
         performTest();
     }
 
+    @Test
     public void testDragIcon_landscape() throws Throwable {
         lockRotation(false);
         performTest();
@@ -58,7 +62,10 @@
 
     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 @@
                 .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());
     }
 }
diff --git a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
index 221fed1..d4d517a 100644
--- a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
@@ -24,10 +24,10 @@
 import android.content.pm.PackageManager;
 import android.database.Cursor;
 import android.os.Bundle;
+import android.support.test.filters.LargeTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.support.test.uiautomator.UiSelector;
-import android.test.suitebuilder.annotation.LargeTest;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppWidgetHost;
 import com.android.launcher3.LauncherAppWidgetHostView;
 import com.android.launcher3.LauncherAppWidgetInfo;
@@ -38,23 +38,40 @@
 import com.android.launcher3.Workspace;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
 import com.android.launcher3.compat.PackageInstallerCompat;
-import com.android.launcher3.ui.LauncherInstrumentationTestCase;
+import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.ContentWriter;
 import com.android.launcher3.util.LooperExecutor;
+import com.android.launcher3.util.rule.LauncherActivityRule;
+import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.WidgetHostViewLoader;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
 /**
  * Tests for bind widget flow.
  *
  * Note running these tests will clear the workspace on the device.
  */
 @LargeTest
-public class BindWidgetTest extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class BindWidgetTest extends AbstractLauncherUiTest {
+
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grandWidgetBind();
 
     private ContentResolver mResolver;
     private AppWidgetManagerCompat mWidgetManager;
@@ -65,21 +82,20 @@
     private int mSessionId = -1;
 
     @Override
-    protected void setUp() throws Exception {
+    @Before
+    public void setUp() throws Exception {
         super.setUp();
 
         mResolver = mTargetContext.getContentResolver();
         mWidgetManager = AppWidgetManagerCompat.getInstance(mTargetContext);
-        grantWidgetPermission();
 
         // Clear all existing data
         LauncherSettings.Settings.call(mResolver, LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
         LauncherSettings.Settings.call(mResolver, LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
+    @After
+    public void tearDown() throws Exception {
         if (mCursor != null) {
             mCursor.close();
         }
@@ -89,6 +105,7 @@
         }
     }
 
+    @Test
     public void testBindNormalWidget_withConfig() {
         LauncherAppWidgetProviderInfo info = findWidgetProvider(true);
         LauncherAppWidgetInfo item = createWidgetInfo(info, true);
@@ -96,6 +113,7 @@
         setupAndVerifyContents(item, LauncherAppWidgetHostView.class, info.label);
     }
 
+    @Test
     public void testBindNormalWidget_withoutConfig() {
         LauncherAppWidgetProviderInfo info = findWidgetProvider(false);
         LauncherAppWidgetInfo item = createWidgetInfo(info, true);
@@ -103,6 +121,7 @@
         setupAndVerifyContents(item, LauncherAppWidgetHostView.class, info.label);
     }
 
+    @Test
     public void testUnboundWidget_removed() throws Exception {
         LauncherAppWidgetProviderInfo info = findWidgetProvider(false);
         LauncherAppWidgetInfo item = createWidgetInfo(info, false);
@@ -121,6 +140,7 @@
         assertFalse(mDevice.findObject(new UiSelector().description(info.label)).exists());
     }
 
+    @Test
     public void testPendingWidget_autoRestored() {
         // A non-restored widget with no config screen gets restored automatically.
         LauncherAppWidgetProviderInfo info = findWidgetProvider(false);
@@ -132,6 +152,7 @@
         setupAndVerifyContents(item, LauncherAppWidgetHostView.class, info.label);
     }
 
+    @Test
     public void testPendingWidget_withConfigScreen() throws Exception {
         // A non-restored widget with config screen get bound and shows a 'Click to setup' UI.
         LauncherAppWidgetProviderInfo info = findWidgetProvider(true);
@@ -154,6 +175,7 @@
                 LauncherSettings.Favorites.APPWIDGET_ID))));
     }
 
+    @Test
     public void testPendingWidget_notRestored_removed() throws Exception {
         LauncherAppWidgetInfo item = getInvalidWidgetInfo();
         item.restoreStatus = LauncherAppWidgetInfo.FLAG_ID_NOT_VALID
@@ -170,6 +192,7 @@
         assertEquals(0, mCursor.getCount());
     }
 
+    @Test
     public void testPendingWidget_notRestored_brokenInstall() throws Exception {
         // A widget which is was being installed once, even if its not being
         // installed at the moment is not removed.
@@ -192,6 +215,7 @@
                         & LauncherAppWidgetInfo.FLAG_ID_NOT_VALID);
     }
 
+    @Test
     public void testPendingWidget_notRestored_activeInstall() throws Exception {
         // A widget which is being installed is not removed
         LauncherAppWidgetInfo item = getInvalidWidgetInfo();
@@ -250,7 +274,7 @@
         resetLoaderState();
 
         // Launch the home activity
-        startLauncher();
+        mActivityMonitor.startLauncher();
         // Verify UI
         UiSelector selector = new UiSelector().packageName(mTargetContext.getPackageName())
                 .className(widgetClass);
diff --git a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
index b798dfa..4b9d83f 100644
--- a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
@@ -15,23 +15,20 @@
  */
 package com.android.launcher3.ui.widget;
 
-import android.app.Activity;
-import android.app.Application;
 import android.app.PendingIntent;
 import android.appwidget.AppWidgetManager;
 import android.content.Intent;
 import android.graphics.Color;
+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.support.test.uiautomator.Until;
-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.LauncherSettings.Favorites;
-import com.android.launcher3.MainThreadExecutor;
 import com.android.launcher3.R;
 import com.android.launcher3.ShortcutInfo;
 import com.android.launcher3.Utilities;
@@ -40,50 +37,48 @@
 import com.android.launcher3.testcomponent.AppWidgetNoConfig;
 import com.android.launcher3.testcomponent.AppWidgetWithConfig;
 import com.android.launcher3.testcomponent.RequestPinItemActivity;
-import com.android.launcher3.ui.LauncherInstrumentationTestCase;
+import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Condition;
-import com.android.launcher3.util.SimpleActivityMonitor;
 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.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.util.UUID;
-import java.util.concurrent.Callable;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Test to verify pin item request flow.
  */
 @LargeTest
-public class RequestPinItemTest  extends LauncherInstrumentationTestCase {
+@RunWith(AndroidJUnit4.class)
+public class RequestPinItemTest  extends AbstractLauncherUiTest {
 
-    private SimpleActivityMonitor mActivityMonitor;
-    private MainThreadExecutor mMainThreadExecutor;
+    @Rule public LauncherActivityRule mActivityMonitor = new LauncherActivityRule();
+    @Rule public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grandWidgetBind();
+    @Rule public ShellCommandRule mDefaultLauncherRule = ShellCommandRule.setDefaultLauncher();
 
     private String mCallbackAction;
     private String mShortcutId;
     private int mAppWidgetId;
 
     @Override
-    protected void setUp() throws Exception {
+    @Before
+    public void setUp() throws Exception {
         super.setUp();
-        grantWidgetPermission();
-        setDefaultLauncher();
-
-        mActivityMonitor = new SimpleActivityMonitor();
-        ((Application) getInstrumentation().getTargetContext().getApplicationContext())
-                .registerActivityLifecycleCallbacks(mActivityMonitor);
-        mMainThreadExecutor = new MainThreadExecutor();
-
         mCallbackAction = UUID.randomUUID().toString();
         mShortcutId = UUID.randomUUID().toString();
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        ((Application) getInstrumentation().getTargetContext().getApplicationContext())
-                .unregisterActivityLifecycleCallbacks(mActivityMonitor);
-        super.tearDown();
-    }
-
+    @Test
     public void testPinWidgetNoConfig() throws Throwable {
         runTest("pinWidgetNoConfig", true, new ItemOperator() {
             @Override
@@ -96,6 +91,7 @@
         });
     }
 
+    @Test
     public void testPinWidgetNoConfig_customPreview() throws Throwable {
         // Command to set custom preview
         Intent command =  RequestPinItemActivity.getCommandIntent(
@@ -113,6 +109,7 @@
         }, command);
     }
 
+    @Test
     public void testPinWidgetWithConfig() throws Throwable {
         runTest("pinWidgetWithConfig", true, new ItemOperator() {
             @Override
@@ -125,6 +122,7 @@
         });
     }
 
+    @Test
     public void testPinShortcut() throws Throwable {
         // Command to set the shortcut id
         Intent command = RequestPinItemActivity.getCommandIntent(
@@ -149,7 +147,7 @@
         lockRotation(true);
 
         clearHomescreen();
-        startLauncher();
+        mActivityMonitor.startLauncher();
 
         // Open all apps and wait for load complete
         final UiObject2 appsContainer = openAllApps();
@@ -191,14 +189,14 @@
         }
 
         // Go back to home
-        mTargetContext.startActivity(getHomeIntent());
+        mActivityMonitor.returnToHome();
         assertTrue(Wait.atMost(new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT));
     }
 
     /**
      * Condition for for an item
      */
-    private class ItemSearchCondition extends Condition implements Callable<Boolean> {
+    private class ItemSearchCondition extends Condition {
 
         private final ItemOperator mOp;
 
@@ -208,22 +206,7 @@
 
         @Override
         public boolean isTrue() throws Throwable {
-            return mMainThreadExecutor.submit(this).get();
-        }
-
-        @Override
-        public Boolean call() throws Exception {
-            // Find the resumed launcher
-            Launcher launcher = null;
-            for (Activity a : mActivityMonitor.resumed) {
-                if (a instanceof Launcher) {
-                    launcher = (Launcher) a;
-                }
-            }
-            if (launcher == null) {
-                return false;
-            }
-            return launcher.getWorkspace().getFirstMatch(mOp) != null;
+            return mMainThreadExecutor.submit(mActivityMonitor.itemExists(mOp)).get();
         }
     }
 }
diff --git a/tests/src/com/android/launcher3/util/SimpleActivityMonitor.java b/tests/src/com/android/launcher3/util/SimpleActivityMonitor.java
deleted file mode 100644
index 6154ab6..0000000
--- a/tests/src/com/android/launcher3/util/SimpleActivityMonitor.java
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Copyright (C) 2017 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License"); you may not
- * use this file except in compliance with the License. You may obtain a copy of
- * the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
- * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
- * License for the specific language governing permissions and limitations under
- * the License.
- */
-package com.android.launcher3.util;
-
-import android.app.Activity;
-import android.app.Application.*;
-import android.os.Bundle;
-
-import java.util.ArrayList;
-
-/**
- * Simple monitor to keep a list of active activities.
- */
-public class SimpleActivityMonitor implements ActivityLifecycleCallbacks {
-
-    public final ArrayList<Activity> created = new ArrayList<>();
-    public final ArrayList<Activity> started = new ArrayList<>();
-    public final ArrayList<Activity> resumed = new ArrayList<>();
-
-    @Override
-    public void onActivityCreated(Activity activity, Bundle bundle) {
-        created.add(activity);
-    }
-
-    @Override
-    public void onActivityStarted(Activity activity) {
-        started.add(activity);
-    }
-
-    @Override
-    public void onActivityResumed(Activity activity) {
-        resumed.add(activity);
-    }
-
-    @Override
-    public void onActivityPaused(Activity activity) {
-        resumed.remove(activity);
-    }
-
-    @Override
-    public void onActivityStopped(Activity activity) {
-        started.remove(activity);
-    }
-
-    @Override
-    public void onActivitySaveInstanceState(Activity activity, Bundle bundle) { }
-
-    @Override
-    public void onActivityDestroyed(Activity activity) {
-        created.remove(activity);
-    }
-}
diff --git a/tests/src/com/android/launcher3/util/rule/LauncherActivityRule.java b/tests/src/com/android/launcher3/util/rule/LauncherActivityRule.java
new file mode 100644
index 0000000..edd152a
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/LauncherActivityRule.java
@@ -0,0 +1,132 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.android.launcher3.util.rule;
+
+import android.app.Activity;
+import android.app.Application;
+import android.app.Application.ActivityLifecycleCallbacks;
+import android.content.Intent;
+import android.os.Bundle;
+import android.support.test.InstrumentationRegistry;
+
+import com.android.launcher3.Launcher;
+import com.android.launcher3.Workspace.ItemOperator;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.util.concurrent.Callable;
+
+/**
+ * Test rule to get the current Launcher activity.
+ */
+public class LauncherActivityRule implements TestRule {
+
+    private Launcher mActivity;
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        return new MyStatement(base);
+    }
+
+    public Launcher getActivity() {
+        return mActivity;
+    }
+
+    public Callable<Boolean> itemExists(final ItemOperator op) {
+        return new Callable<Boolean>() {
+
+            @Override
+            public Boolean call() throws Exception {
+                Launcher launcher = getActivity();
+                if (launcher == null) {
+                    return false;
+                }
+                return launcher.getWorkspace().getFirstMatch(op) != null;
+            }
+        };
+    }
+
+    /**
+     * Starts the launcher activity in the target package.
+     */
+    public void startLauncher() {
+        InstrumentationRegistry.getInstrumentation().startActivitySync(getHomeIntent());
+    }
+
+    public void returnToHome() {
+        InstrumentationRegistry.getTargetContext().startActivity(getHomeIntent());
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    public static Intent getHomeIntent() {
+        return new Intent(Intent.ACTION_MAIN)
+                .addCategory(Intent.CATEGORY_HOME)
+                .setPackage(InstrumentationRegistry.getTargetContext().getPackageName())
+                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
+    }
+
+    private class MyStatement extends Statement implements ActivityLifecycleCallbacks {
+
+        private final Statement mBase;
+
+        public MyStatement(Statement base) {
+            mBase = base;
+        }
+
+        @Override
+        public void evaluate() throws Throwable {
+            Application app = (Application)
+                    InstrumentationRegistry.getTargetContext().getApplicationContext();
+            app.registerActivityLifecycleCallbacks(this);
+            try {
+                mBase.evaluate();
+            } finally {
+                app.unregisterActivityLifecycleCallbacks(this);
+            }
+        }
+
+        @Override
+        public void onActivityCreated(Activity activity, Bundle bundle) {
+            if (activity instanceof Launcher) {
+                mActivity = (Launcher) activity;
+            }
+        }
+
+        @Override
+        public void onActivityStarted(Activity activity) { }
+
+        @Override
+        public void onActivityResumed(Activity activity) { }
+
+        @Override
+        public void onActivityPaused(Activity activity) { }
+
+        @Override
+        public void onActivityStopped(Activity activity) { }
+
+        @Override
+        public void onActivitySaveInstanceState(Activity activity, Bundle bundle) { }
+
+        @Override
+        public void onActivityDestroyed(Activity activity) {
+            if (activity == mActivity) {
+                mActivity = null;
+            }
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
new file mode 100644
index 0000000..dba2d71
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.android.launcher3.util.rule;
+
+import android.content.ComponentName;
+import android.content.pm.ActivityInfo;
+import android.os.ParcelFileDescriptor;
+import android.support.test.InstrumentationRegistry;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.io.FileInputStream;
+import java.io.IOException;
+
+/**
+ * Test rule which executes a shell command at the start of the test.
+ */
+public class ShellCommandRule implements TestRule {
+
+    private final String mCmd;
+
+    public ShellCommandRule(String cmd) {
+        mCmd = cmd;
+    }
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        return new MyStatement(base, mCmd);
+    }
+
+    public static void runShellCommand(String command) throws IOException {
+        ParcelFileDescriptor pfd = InstrumentationRegistry.getInstrumentation().getUiAutomation()
+                .executeShellCommand(command);
+
+        // Read the input stream fully.
+        FileInputStream fis = new ParcelFileDescriptor.AutoCloseInputStream(pfd);
+        while (fis.read() != -1);
+        fis.close();
+    }
+
+    private static class MyStatement extends Statement {
+        private final Statement mBase;
+        private final String mCmd;
+
+        public MyStatement(Statement base, String cmd) {
+            mBase = base;
+            mCmd = cmd;
+        }
+
+        @Override
+        public void evaluate() throws Throwable {
+            runShellCommand(mCmd);
+            mBase.evaluate();
+        }
+    }
+
+    /**
+     * Grants the launcher permission to bind widgets.
+     */
+    public static ShellCommandRule grandWidgetBind() {
+        return new ShellCommandRule("appwidget grantbind --package "
+                + InstrumentationRegistry.getTargetContext().getPackageName());
+    }
+
+    /**
+     * Sets the target launcher as default launcher.
+     */
+    public static ShellCommandRule setDefaultLauncher() {
+        ActivityInfo launcher = InstrumentationRegistry.getTargetContext().getPackageManager()
+                .queryIntentActivities(LauncherActivityRule.getHomeIntent(), 0).get(0)
+                .activityInfo;
+        return new ShellCommandRule("cmd package set-home-activity " +
+                new ComponentName(launcher.packageName, launcher.name).flattenToString());
+    }
+}