diff --git a/tests/Android.mk b/tests/Android.mk
index 466146e..5103ced 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -17,9 +17,9 @@
 include $(CLEAR_VARS)
 
 LOCAL_MODULE_TAGS := tests
-#LOCAL_STATIC_JAVA_LIBRARIES := android-support-test ub-uiautomator mockito-target-minus-junit4
+LOCAL_STATIC_JAVA_LIBRARIES := android-support-test ub-uiautomator mockito-target-minus-junit4
 
-#LOCAL_SRC_FILES := $(call all-java-files-under, src)
+LOCAL_SRC_FILES := $(call all-java-files-under, src)
 
 LOCAL_SDK_VERSION := current
 LOCAL_MIN_SDK_VERSION := 21
diff --git a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
index f23a574..1d092ab 100644
--- a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
+++ b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
@@ -8,6 +8,7 @@
 import android.graphics.Bitmap.Config;
 import android.os.Process;
 import android.os.UserHandle;
+import android.support.annotation.NonNull;
 import android.support.test.InstrumentationRegistry;
 import android.test.ProviderTestCase2;
 
@@ -24,6 +25,7 @@
 import com.android.launcher3.compat.LauncherActivityInfoCompat;
 import com.android.launcher3.config.ProviderConfig;
 import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.Provider;
 import com.android.launcher3.util.TestLauncherProvider;
 
 import org.mockito.ArgumentCaptor;
@@ -184,9 +186,10 @@
         }
 
         @Override
-        protected CacheEntry cacheLocked(ComponentName componentName,
-                LauncherActivityInfoCompat info, UserHandle user,
-                boolean usePackageIcon, boolean useLowResIcon) {
+        protected CacheEntry cacheLocked(
+                @NonNull ComponentName componentName,
+                @NonNull Provider<LauncherActivityInfoCompat> infoProvider,
+                UserHandle user, boolean usePackageIcon, boolean useLowResIcon) {
             CacheEntry entry = mCache.get(new ComponentKey(componentName, user));
             if (entry == null) {
                 entry = new CacheEntry();