diff --git a/build.gradle b/build.gradle
index f357e5c..61c05e5 100644
--- a/build.gradle
+++ b/build.gradle
@@ -80,18 +80,18 @@
         }
 
         aosp {
-            java.srcDirs = ['src_flags']
+            java.srcDirs = ['src_flags', "src_ui_overrides"]
         }
 
         l3go {
             res.srcDirs = ['go/res']
-            java.srcDirs = ['go/src_flags']
+            java.srcDirs = ['go/src_flags', "src_ui_overrides"]
             manifest.srcFile "go/AndroidManifest.xml"
         }
 
         quickstep {
             res.srcDirs = ['quickstep/res']
-            java.srcDirs = ['quickstep/src_flags', 'quickstep/src']
+            java.srcDirs = ['src_flags', 'quickstep/src']
             manifest.srcFile "quickstep/AndroidManifest.xml"
         }
     }
@@ -115,7 +115,8 @@
     androidTestCompile "org.mockito:mockito-core:1.9.5"
     androidTestCompile 'com.google.dexmaker:dexmaker:1.2'
     androidTestCompile 'com.google.dexmaker:dexmaker-mockito:1.2'
-    androidTestCompile 'com.android.support.test:runner:0.5'
+    androidTestCompile 'com.android.support.test:runner:1.0.0'
+    androidTestCompile 'com.android.support.test:rules:1.0.0'
     androidTestCompile 'com.android.support.test.uiautomator:uiautomator-v18:2.1.2'
     androidTestCompile "com.android.support:support-annotations:${SUPPORT_LIBS_VERSION}"
 }
diff --git a/tests/src/com/android/launcher3/logging/FileLogTest.java b/tests/src/com/android/launcher3/logging/FileLogTest.java
index 7048c28..9c7cb8f 100644
--- a/tests/src/com/android/launcher3/logging/FileLogTest.java
+++ b/tests/src/com/android/launcher3/logging/FileLogTest.java
@@ -1,41 +1,51 @@
 package com.android.launcher3.logging;
 
-import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 
 import java.io.File;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.Calendar;
 
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
 /**
  * Tests for {@link FileLog}
  */
 @SmallTest
-public class FileLogTest extends AndroidTestCase {
+@RunWith(AndroidJUnit4.class)
+public class FileLogTest {
 
     private File mTempDir;
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void setUp() throws Exception {
         int count = 0;
         do {
-            mTempDir = new File(getContext().getCacheDir(), "log-test-" + (count++));
+            mTempDir = new File(InstrumentationRegistry.getTargetContext().getCacheDir(),
+                    "log-test-" + (count++));
         } while(!mTempDir.mkdir());
 
         FileLog.setDir(mTempDir);
     }
 
-    @Override
-    protected void tearDown() throws Exception {
+    @After
+    public void tearDown() throws Exception {
         // Clear existing logs
         new File(mTempDir, "log-0").delete();
         new File(mTempDir, "log-1").delete();
         mTempDir.delete();
-        super.tearDown();
     }
 
+    @Test
     public void testPrintLog() throws Exception {
         if (!FileLog.ENABLED) {
             return;
@@ -56,6 +66,7 @@
         assertTrue(writer.toString().contains("hoolalala"));
     }
 
+    @Test
     public void testOldFileTruncated() throws Exception {
         if (!FileLog.ENABLED) {
             return;
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
index a486ceb..401711d 100644
--- a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -7,6 +7,7 @@
 import android.content.Intent;
 import android.graphics.Rect;
 import android.net.Uri;
+import android.support.test.runner.AndroidJUnit4;
 import android.util.Pair;
 
 import com.android.launcher3.ItemInfo;
@@ -15,21 +16,25 @@
 import com.android.launcher3.ShortcutInfo;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.LongArrayMap;
-import com.android.launcher3.util.Provider;
 
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 
 import java.util.ArrayList;
 import java.util.List;
 
-import static org.mockito.Matchers.isNull;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 /**
  * Tests for {@link AddWorkspaceItemsTask}
  */
+@RunWith(AndroidJUnit4.class)
 public class AddWorkspaceItemsTaskTest extends BaseModelUpdateTaskTestCase {
 
     private final ComponentName mComponent1 = new ComponentName("a", "b");
@@ -39,9 +44,8 @@
     private ArrayList<Long> newScreens;
     private LongArrayMap<GridOccupancy> screenOccupancy;
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void initData() throws Exception {
         existingScreens = new ArrayList<>();
         screenOccupancy = new LongArrayMap<>();
         newScreens = new ArrayList<>();
@@ -62,6 +66,7 @@
         };
     }
 
+    @Test
     public void testFindSpaceForItem_prefers_second() {
         // First screen has only one hole of size 1
         int nextId = setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
@@ -83,13 +88,12 @@
                 .isRegionVacant(spaceFound.second[0], spaceFound.second[1], 2, 3));
     }
 
+    @Test
     public void testFindSpaceForItem_adds_new_screen() throws Exception {
         // First screen has 2 holes of sizes 3x2 and 2x3
         setupWorkspaceWithHoles(1, 1, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
         commitScreensToDb();
 
-        when(appState.getContext()).thenReturn(getMockContext());
-
         ArrayList<Long> oldScreens = new ArrayList<>(existingScreens);
         Pair<Long, int[]> spaceFound = newTask()
                 .findSpaceForItem(appState, bgDataModel, existingScreens, newScreens, 3, 3);
@@ -97,6 +101,7 @@
         assertTrue(newScreens.contains(spaceFound.first));
     }
 
+    @Test
     public void testAddItem_existing_item_ignored() throws Exception {
         ShortcutInfo info = new ShortcutInfo();
         info.intent = new Intent().setComponent(mComponent1);
@@ -105,12 +110,11 @@
         setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
         commitScreensToDb();
 
-        when(appState.getContext()).thenReturn(getMockContext());
-
         // Nothing was added
         assertTrue(executeTaskForTest(newTask(info)).isEmpty());
     }
 
+    @Test
     public void testAddItem_some_items_added() throws Exception {
         ShortcutInfo info = new ShortcutInfo();
         info.intent = new Intent().setComponent(mComponent1);
@@ -122,8 +126,6 @@
         setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
         commitScreensToDb();
 
-        when(appState.getContext()).thenReturn(getMockContext());
-
         executeTaskForTest(newTask(info, info2)).get(0).run();
         ArgumentCaptor<ArrayList> notAnimated = ArgumentCaptor.forClass(ArrayList.class);
         ArgumentCaptor<ArrayList> animated = ArgumentCaptor.forClass(ArrayList.class);
@@ -168,7 +170,7 @@
     }
 
     private void commitScreensToDb() throws Exception {
-        LauncherSettings.Settings.call(getMockContentResolver(),
+        LauncherSettings.Settings.call(mProviderRule.getResolver(),
                 LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
 
         Uri uri = LauncherSettings.WorkspaceScreens.CONTENT_URI;
@@ -183,6 +185,6 @@
             v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
             ops.add(ContentProviderOperation.newInsert(uri).withValues(v).build());
         }
-        getMockContentResolver().applyBatch(LauncherProvider.AUTHORITY, ops);
+        mProviderRule.getResolver().applyBatch(LauncherProvider.AUTHORITY, ops);
     }
 }
diff --git a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
index 3d03507..bbb6772 100644
--- a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
+++ b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
@@ -1,7 +1,9 @@
 package com.android.launcher3.model;
 
 import android.content.ComponentName;
+import android.content.ContentResolver;
 import android.content.Context;
+import android.content.ContextWrapper;
 import android.content.Intent;
 import android.content.pm.LauncherActivityInfo;
 import android.content.res.Resources;
@@ -11,7 +13,8 @@
 import android.os.UserHandle;
 import android.support.annotation.NonNull;
 import android.support.test.InstrumentationRegistry;
-import android.test.ProviderTestCase2;
+import android.support.test.rule.provider.ProviderTestRule;
+import android.support.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.AllAppsList;
 import com.android.launcher3.AppFilter;
@@ -28,6 +31,8 @@
 import com.android.launcher3.util.Provider;
 import com.android.launcher3.util.TestLauncherProvider;
 
+import org.junit.Before;
+import org.junit.Rule;
 import org.mockito.ArgumentCaptor;
 
 import java.io.BufferedReader;
@@ -46,7 +51,12 @@
 /**
  * Base class for writing tests for Model update tasks.
  */
-public class BaseModelUpdateTaskTestCase extends ProviderTestCase2<TestLauncherProvider> {
+public class BaseModelUpdateTaskTestCase {
+
+    @Rule
+    public ProviderTestRule mProviderRule =
+            new ProviderTestRule.Builder(TestLauncherProvider.class, LauncherProvider.AUTHORITY)
+                    .build();
 
     public final HashMap<Class, HashMap<String, Field>> fieldCache = new HashMap<>();
 
@@ -63,14 +73,8 @@
     public AllAppsList allAppsList;
     public Callbacks callbacks;
 
-    public BaseModelUpdateTaskTestCase() {
-        super(TestLauncherProvider.class, LauncherProvider.AUTHORITY);
-    }
-
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
+    @Before
+    public void setUp() throws Exception {
         callbacks = mock(Callbacks.class);
         appState = mock(LauncherAppState.class);
         model = mock(LauncherModel.class);
@@ -83,7 +87,12 @@
         myUser = Process.myUserHandle();
 
         bgDataModel = new BgDataModel();
-        targetContext = InstrumentationRegistry.getTargetContext();
+        targetContext = new ContextWrapper(InstrumentationRegistry.getTargetContext()) {
+            @Override
+            public ContentResolver getContentResolver() {
+                return mProviderRule.getResolver();
+            }
+        };
         idp = new InvariantDeviceProfile();
         iconCache = new MyIconCache(targetContext, idp);
 
@@ -91,6 +100,8 @@
 
         when(appState.getIconCache()).thenReturn(iconCache);
         when(appState.getInvariantDeviceProfile()).thenReturn(idp);
+        when(appState.getContext()).thenReturn(targetContext);
+
     }
 
     /**
diff --git a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index d595e6c..ac9d319 100644
--- a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -1,23 +1,34 @@
 package com.android.launcher3.model;
 
+import android.support.test.runner.AndroidJUnit4;
+
 import com.android.launcher3.AppInfo;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.ShortcutInfo;
 
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.util.Arrays;
 import java.util.HashSet;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+
 /**
  * Tests for {@link CacheDataUpdatedTask}
  */
+@RunWith(AndroidJUnit4.class)
 public class CacheDataUpdatedTaskTest extends BaseModelUpdateTaskTestCase {
 
     private static final String NEW_LABEL_PREFIX = "new-label-";
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-
+    @Before
+    public void initData() throws Exception {
         initializeData("cache_data_updated_task_data");
         // Add dummy entries in the cache to simulate update
         for (ItemInfo info : bgDataModel.itemsIdMap) {
@@ -29,6 +40,7 @@
         return new CacheDataUpdatedTask(op, myUser, new HashSet<>(Arrays.asList(pkg)));
     }
 
+    @Test
     public void testCacheUpdate_update_apps() throws Exception {
         // Clear all icons from apps list so that its easy to check what was updated
         for (AppInfo info : allAppsList.data) {
@@ -52,6 +64,7 @@
         }
     }
 
+    @Test
     public void testSessionUpdate_ignores_normal_apps() throws Exception {
         executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app1"));
 
@@ -59,6 +72,7 @@
         verifyUpdate();
     }
 
+    @Test
     public void testSessionUpdate_updates_pending_apps() throws Exception {
         executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app3"));
 
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
index fd62d36..b92f612 100644
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
@@ -1,11 +1,16 @@
 package com.android.launcher3.model;
 
+import android.content.ContentResolver;
 import android.content.ContentValues;
+import android.content.Context;
+import android.content.ContextWrapper;
 import android.content.Intent;
 import android.database.Cursor;
 import android.graphics.Point;
-import android.test.ProviderTestCase2;
-import android.test.suitebuilder.annotation.MediumTest;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.MediumTest;
+import android.support.test.rule.provider.ProviderTestRule;
+import android.support.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherModel;
@@ -15,15 +20,29 @@
 import com.android.launcher3.model.GridSizeMigrationTask.MultiStepMigrationTask;
 import com.android.launcher3.util.TestLauncherProvider;
 
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.LinkedList;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
 /**
  * Unit tests for {@link GridSizeMigrationTask}
  */
 @MediumTest
-public class GridSizeMigrationTaskTest extends ProviderTestCase2<TestLauncherProvider> {
+@RunWith(AndroidJUnit4.class)
+public class GridSizeMigrationTaskTest {
+
+    @Rule
+    public ProviderTestRule mProviderRule =
+            new ProviderTestRule.Builder(TestLauncherProvider.class, LauncherProvider.AUTHORITY)
+                    .build();
 
     private static final long DESKTOP = LauncherSettings.Favorites.CONTAINER_DESKTOP;
     private static final long HOTSEAT = LauncherSettings.Favorites.CONTAINER_HOTSEAT;
@@ -37,20 +56,25 @@
 
     private HashSet<String> mValidPackages;
     private InvariantDeviceProfile mIdp;
+    private Context mContext;
 
-    public GridSizeMigrationTaskTest() {
-        super(TestLauncherProvider.class, LauncherProvider.AUTHORITY);
-    }
-
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void setUp() throws Exception {
         mValidPackages = new HashSet<>();
         mValidPackages.add(TEST_PACKAGE);
 
         mIdp = new InvariantDeviceProfile();
+
+        mContext = new ContextWrapper(InstrumentationRegistry.getTargetContext()) {
+
+            @Override
+            public ContentResolver getContentResolver() {
+                return mProviderRule.getResolver();
+            }
+        };
     }
 
+    @Test
     public void testHotseatMigration_apps_dropped() throws Exception {
         long[] hotseatItems = {
                 addItem(APPLICATION, 0, HOTSEAT, 0, 0),
@@ -61,7 +85,7 @@
         };
 
         mIdp.numHotseatIcons = 3;
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, 5, 3)
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages, 5, 3)
                 .migrateHotseat();
         if (FeatureFlags.NO_ALL_APPS_ICON) {
             // First item is dropped as it has the least weight.
@@ -72,6 +96,7 @@
         }
     }
 
+    @Test
     public void testHotseatMigration_shortcuts_dropped() throws Exception {
         long[] hotseatItems = {
                 addItem(APPLICATION, 0, HOTSEAT, 0, 0),
@@ -82,7 +107,7 @@
         };
 
         mIdp.numHotseatIcons = 3;
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, 5, 3)
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages, 5, 3)
                 .migrateHotseat();
         if (FeatureFlags.NO_ALL_APPS_ICON) {
             // First item is dropped as it has the least weight.
@@ -98,7 +123,7 @@
         int total = 0;
 
         for (long id : sortedIds) {
-            Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+            Cursor c = mProviderRule.getResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                     new String[]{LauncherSettings.Favorites._ID},
                     "container=-101 and screen=" + screenId, null, null, null);
 
@@ -116,13 +141,14 @@
         }
 
         // Verify that not other entry exist in the DB.
-        Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+        Cursor c = mProviderRule.getResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites._ID},
                 "container=-101", null, null, null);
         assertEquals(total, c.getCount());
         c.close();
     }
 
+    @Test
     public void testWorkspace_empty_row_column_removed() throws Exception {
         long[][][] ids = createGrid(new int[][][]{{
                 {  0,  0, -1,  1},
@@ -131,7 +157,7 @@
                 {  5,  2, -1,  6},
         }});
 
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages,
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages,
                 new Point(4, 4), new Point(3, 3)).migrateWorkspace();
 
         // Column 2 and row 2 got removed.
@@ -142,6 +168,7 @@
         }});
     }
 
+    @Test
     public void testWorkspace_new_screen_created() throws Exception {
         long[][][] ids = createGrid(new int[][][]{{
                 {  0,  0,  0,  1},
@@ -150,7 +177,7 @@
                 {  5,  2, -1,  6},
         }});
 
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages,
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages,
                 new Point(4, 4), new Point(3, 3)).migrateWorkspace();
 
         // Items in the second column get moved to new screen
@@ -163,6 +190,7 @@
         }});
     }
 
+    @Test
     public void testWorkspace_items_merged_in_next_screen() throws Exception {
         long[][][] ids = createGrid(new int[][][]{{
                 {  0,  0,  0,  1},
@@ -174,7 +202,7 @@
                 {  3,  1, -1,  4},
         }});
 
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages,
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages,
                 new Point(4, 4), new Point(3, 3)).migrateWorkspace();
 
         // Items in the second column of the first screen should get placed on the 3rd
@@ -190,6 +218,7 @@
         }});
     }
 
+    @Test
     public void testWorkspace_items_not_merged_in_next_screen() throws Exception {
         // First screen has 2 items that need to be moved, but second screen has only one
         // empty space after migration (top-left corner)
@@ -205,7 +234,7 @@
                 {  5,  2, -1,  6},
         }});
 
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages,
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages,
                 new Point(4, 4), new Point(3, 3)).migrateWorkspace();
 
         // Items in the second column of the first screen should get placed on a new screen.
@@ -222,6 +251,7 @@
         }});
     }
 
+    @Test
     public void testWorkspace_first_row_blocked() throws Exception {
         // The first screen has one item on the 4th column which needs moving, as the first row
         // will be kept empty.
@@ -232,7 +262,7 @@
                 {  5,  2,  7, -1},
         }}, 0);
 
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages,
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages,
                 new Point(4, 4), new Point(3, 4)).migrateWorkspace();
 
         // Items in the second column of the first screen should get placed on a new screen.
@@ -246,6 +276,7 @@
         }});
     }
 
+    @Test
     public void testWorkspace_items_moved_to_empty_first_row() throws Exception {
         // Items will get moved to the next screen to keep the first screen empty.
         long[][][] ids = createGrid(new int[][][]{{
@@ -255,7 +286,7 @@
                 {  5,  6,  7, -1},
         }}, 0);
 
-        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages,
+        new GridSizeMigrationTask(mContext, mIdp, mValidPackages,
                 new Point(4, 4), new Point(3, 3)).migrateWorkspace();
 
         // Items in the second column of the first screen should get placed on a new screen.
@@ -281,7 +312,7 @@
      * @return the same grid representation where each entry is the corresponding item id.
      */
     private long[][][] createGrid(int[][][] typeArray, long startScreen) throws Exception {
-        LauncherSettings.Settings.call(getMockContentResolver(),
+        LauncherSettings.Settings.call(mProviderRule.getResolver(),
                 LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
         long[][][] ids = new long[typeArray.length][][];
 
@@ -290,13 +321,13 @@
             long screenId = startScreen + i;
 
             // Keep the screen id counter up to date
-            LauncherSettings.Settings.call(getMockContentResolver(),
+            LauncherSettings.Settings.call(mProviderRule.getResolver(),
                     LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
 
             ContentValues v = new ContentValues();
             v.put(LauncherSettings.WorkspaceScreens._ID, screenId);
             v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
-            getMockContentResolver().insert(LauncherSettings.WorkspaceScreens.CONTENT_URI, v);
+            mProviderRule.getResolver().insert(LauncherSettings.WorkspaceScreens.CONTENT_URI, v);
 
             ids[i] = new long[typeArray[i].length][];
             for (int y = 0; y < typeArray[i].length; y++) {
@@ -320,7 +351,7 @@
      *            represent the workspace grid.
      */
     private void verifyWorkspace(long[][][] ids) {
-        ArrayList<Long> allScreens = LauncherModel.loadWorkspaceScreensDb(getMockContext());
+        ArrayList<Long> allScreens = LauncherModel.loadWorkspaceScreensDb(mContext);
         assertEquals(ids.length, allScreens.size());
         int total = 0;
 
@@ -330,7 +361,8 @@
                 for (int x = 0; x < ids[i][y].length; x++) {
                     long id = ids[i][y][x];
 
-                    Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                    Cursor c = mProviderRule.getResolver().query(
+                            LauncherSettings.Favorites.CONTENT_URI,
                             new String[]{LauncherSettings.Favorites._ID},
                             "container=-100 and screen=" + screenId +
                                     " and cellX=" + x + " and cellY=" + y, null, null, null);
@@ -349,7 +381,7 @@
         }
 
         // Verify that not other entry exist in the DB.
-        Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+        Cursor c = mProviderRule.getResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites._ID},
                 "container=-100", null, null, null);
         assertEquals(total, c.getCount());
@@ -362,7 +394,7 @@
      *             folder (where the type represents the number of items in the folder).
      */
     private long addItem(int type, long screen, long container, int x, int y) throws Exception {
-        long id = LauncherSettings.Settings.call(getMockContentResolver(),
+        long id = LauncherSettings.Settings.call(mProviderRule.getResolver(),
                 LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
                 .getLong(LauncherSettings.Settings.EXTRA_VALUE);
 
@@ -387,16 +419,18 @@
             }
         }
 
-        getMockContentResolver().insert(LauncherSettings.Favorites.CONTENT_URI, values);
+        mProviderRule.getResolver().insert(LauncherSettings.Favorites.CONTENT_URI, values);
         return id;
     }
 
+    @Test
     public void testMultiStepMigration_small_to_large() throws Exception {
         MultiStepMigrationTaskVerifier verifier = new MultiStepMigrationTaskVerifier();
         verifier.migrate(new Point(3, 3), new Point(5, 5));
         verifier.assertCompleted();
     }
 
+    @Test
     public void testMultiStepMigration_large_to_small() throws Exception {
         MultiStepMigrationTaskVerifier verifier = new MultiStepMigrationTaskVerifier(
                 5, 5, 4, 4,
@@ -406,6 +440,7 @@
         verifier.assertCompleted();
     }
 
+    @Test
     public void testMultiStepMigration_zig_zag() throws Exception {
         MultiStepMigrationTaskVerifier verifier = new MultiStepMigrationTaskVerifier(
                 5, 7, 4, 7,
diff --git a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
index ed893c4..0a741c4 100644
--- a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
+++ b/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -1,22 +1,30 @@
 package com.android.launcher3.model;
 
+import android.support.test.runner.AndroidJUnit4;
+
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.ShortcutInfo;
 import com.android.launcher3.compat.PackageInstallerCompat;
 import com.android.launcher3.compat.PackageInstallerCompat.PackageInstallInfo;
 
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.util.Arrays;
 import java.util.HashSet;
 
+import static org.junit.Assert.assertEquals;
+
 /**
  * Tests for {@link PackageInstallStateChangedTask}
  */
+@RunWith(AndroidJUnit4.class)
 public class PackageInstallStateChangedTaskTest extends BaseModelUpdateTaskTestCase {
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void initData() throws Exception {
         initializeData("package_install_state_change_task_data");
     }
 
@@ -26,6 +34,7 @@
         return new PackageInstallStateChangedTask(installInfo);
     }
 
+    @Test
     public void testSessionUpdate_ignore_installed() throws Exception {
         executeTaskForTest(newTask("app1", 30));
 
@@ -33,12 +42,14 @@
         verifyProgressUpdate(0);
     }
 
+    @Test
     public void testSessionUpdate_shortcuts_updated() throws Exception {
         executeTaskForTest(newTask("app3", 30));
 
         verifyProgressUpdate(30, 5L, 6L, 7L);
     }
 
+    @Test
     public void testSessionUpdate_widgets_updated() throws Exception {
         executeTaskForTest(newTask("app4", 30));
 
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 5858e13..5d417b5 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -3,23 +3,32 @@
 import android.content.ContentValues;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
-import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.MediumTest;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.MediumTest;
+import android.support.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
 
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+
 /**
  * Tests for {@link RestoreDbTask}
  */
 @MediumTest
-public class RestoreDbTaskTest extends AndroidTestCase {
+@RunWith(AndroidJUnit4.class)
+public class RestoreDbTaskTest {
 
+    @Test
     public void testGetProfileId() throws Exception {
         SQLiteDatabase db = new MyDatabaseHelper(23).getWritableDatabase();
         assertEquals(23, new RestoreDbTask().getDefaultProfileId(db));
     }
 
+    @Test
     public void testMigrateProfileId() throws Exception {
         SQLiteDatabase db = new MyDatabaseHelper(42).getWritableDatabase();
         // Add some dummy data
@@ -57,7 +66,7 @@
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(getContext(), null, null);
+            super(InstrumentationRegistry.getContext(), null, null);
             mProfileId = profileId;
         }
 
@@ -66,6 +75,9 @@
             return mProfileId;
         }
 
+        @Override
+        protected void handleOneTimeDataUpgrade(SQLiteDatabase db) { }
+
         protected void onEmptyDbCreated() { }
     }
 }
diff --git a/tests/src/com/android/launcher3/util/FocusLogicTest.java b/tests/src/com/android/launcher3/util/FocusLogicTest.java
index 79aed80..691d9bc 100644
--- a/tests/src/com/android/launcher3/util/FocusLogicTest.java
+++ b/tests/src/com/android/launcher3/util/FocusLogicTest.java
@@ -16,30 +16,24 @@
 
 package com.android.launcher3.util;
 
-import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.view.KeyEvent;
-import android.view.View;
 
-import com.android.launcher3.util.FocusLogic;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Tests the {@link FocusLogic} class that handles key event based focus handling.
  */
 @SmallTest
-public final class FocusLogicTest extends AndroidTestCase {
+@RunWith(AndroidJUnit4.class)
+public final class FocusLogicTest {
 
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        // Nothing to set up as this class only tests static methods.
-    }
-
-    @Override
-    protected void tearDown() throws Exception {
-        // Nothing to tear down as this class only tests static methods.
-    }
-
+    @Test
     public void testShouldConsume() {
          assertTrue(FocusLogic.shouldConsume(KeyEvent.KEYCODE_DPAD_LEFT));
          assertTrue(FocusLogic.shouldConsume(KeyEvent.KEYCODE_DPAD_RIGHT));
@@ -51,12 +45,14 @@
          assertTrue(FocusLogic.shouldConsume(KeyEvent.KEYCODE_PAGE_DOWN));
     }
 
+    @Test
     public void testCreateSparseMatrix() {
          // Either, 1) create a helper method to generate/instantiate all possible cell layout that
          // may get created in real world to test this method. OR 2) Move all the matrix
          // management routine to celllayout and write tests for them.
     }
 
+    @Test
     public void testMoveFromBottomRightToBottomLeft() {
         int[][] map = transpose(new int[][] {
                 {-1, 0, -1, -1, -1, -1},
@@ -69,6 +65,7 @@
         assertEquals(1, i);
     }
 
+    @Test
     public void testMoveFromBottomRightToTopLeft() {
         int[][] map = transpose(new int[][] {
                 {-1, 0, -1, -1, -1, -1},
@@ -81,6 +78,7 @@
         assertEquals(FocusLogic.NEXT_PAGE_FIRST_ITEM, i);
     }
 
+    @Test
     public void testMoveIntoHotseatWithEqualHotseatAndWorkspaceColumns() {
         // Test going from an icon right above the All Apps button to the All Apps button.
         int[][] map = transpose(new int[][] {
@@ -105,6 +103,7 @@
         assertEquals(4, i);
     }
 
+    @Test
     public void testMoveIntoHotseatWithExtraColumnForAllApps() {
         // Test going from an icon above and to the left
         // of the All Apps button to the All Apps button.
@@ -187,6 +186,7 @@
         assertEquals(12, i);
     }
 
+    @Test
     public void testCrossingAllAppsColumn() {
         // Test crossing from left to right in portrait.
         int[][] map = transpose(new int[][] {
diff --git a/tests/src/com/android/launcher3/util/GridOccupancyTest.java b/tests/src/com/android/launcher3/util/GridOccupancyTest.java
index 7d0fe71..fdd8e88 100644
--- a/tests/src/com/android/launcher3/util/GridOccupancyTest.java
+++ b/tests/src/com/android/launcher3/util/GridOccupancyTest.java
@@ -1,15 +1,23 @@
 package com.android.launcher3.util;
 
-import android.test.suitebuilder.annotation.SmallTest;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
 
-import junit.framework.TestCase;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Unit tests for {@link GridOccupancy}
  */
 @SmallTest
-public class GridOccupancyTest extends TestCase {
+@RunWith(AndroidJUnit4.class)
+public class GridOccupancyTest {
 
+    @Test
     public void testFindVacantCell() {
         GridOccupancy grid = initGrid(4,
                 1, 1, 1, 0, 0,
@@ -30,6 +38,7 @@
         assertFalse(grid.findVacantCell(vacant, 3, 3));
     }
 
+    @Test
     public void testIsRegionVacant() {
         GridOccupancy grid = initGrid(4,
                 1, 1, 1, 0, 0,