1 files changed, 7 insertions, 12 deletions
diff --git a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
index 13e09869f..3d03507ea 100644
--- a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
+++ b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
@@ -21,7 +21,7 @@ import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
-import com.android.launcher3.LauncherModel.BaseModelUpdateTask;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherModel.Callbacks;
 import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.util.ComponentKey;
@@ -75,8 +75,10 @@ public class BaseModelUpdateTaskTestCase extends ProviderTestCase2<TestLauncherP
         appState = mock(LauncherAppState.class);
         model = mock(LauncherModel.class);
         modelWriter = mock(ModelWriter.class);
+
         when(appState.getModel()).thenReturn(model);
         when(model.getWriter(anyBoolean())).thenReturn(modelWriter);
+        when(model.getCallback()).thenReturn(callbacks);
 
         myUser = Process.myUserHandle();
 
@@ -94,20 +96,13 @@ public class BaseModelUpdateTaskTestCase extends ProviderTestCase2<TestLauncherP
     /**
      * Synchronously executes the task and returns all the UI callbacks posted.
      */
-    public List<Runnable> executeTaskForTest(BaseModelUpdateTask task) throws Exception {
-        LauncherModel mockModel = mock(LauncherModel.class);
-        when(mockModel.getCallback()).thenReturn(callbacks);
-
-        Field f = BaseModelUpdateTask.class.getDeclaredField("mModel");
-        f.setAccessible(true);
-        f.set(task, mockModel);
+    public List<Runnable> executeTaskForTest(ModelUpdateTask task) throws Exception {
+        when(model.isModelLoaded()).thenReturn(true);
 
         Executor mockExecutor = mock(Executor.class);
-        f = BaseModelUpdateTask.class.getDeclaredField("mUiExecutor");
-        f.setAccessible(true);
-        f.set(task, mockExecutor);
 
-        task.execute(appState, bgDataModel, allAppsList);
+        task.init(appState, model, bgDataModel, allAppsList, mockExecutor);
+        task.run();
         ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
         verify(mockExecutor, atLeast(0)).execute(captor.capture());