diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 48c5c56..82bee0e 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -42,7 +42,7 @@
 import com.android.launcher3.model.AddWorkspaceItemsTask;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.CacheDataUpdatedTask;
-import com.android.launcher3.model.ExtendedModelTask;
+import com.android.launcher3.model.BaseModelUpdateTask;
 import com.android.launcher3.model.LoaderResults;
 import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelWriter;
@@ -80,7 +80,6 @@
  */
 public class LauncherModel extends BroadcastReceiver
         implements LauncherAppsCompat.OnAppsChangedCallbackCompat {
-    static final boolean DEBUG_TASKS = false;
     private static final boolean DEBUG_RECEIVER = false;
 
     static final String TAG = "Launcher.Model";
@@ -425,7 +424,7 @@
     public void forceReload() {
         synchronized (mLock) {
             // Stop any existing loaders first, so they don't set mModelLoaded to true later
-            stopLoaderLocked();
+            stopLoader();
             mModelLoaded = false;
         }
 
@@ -451,16 +450,6 @@
         }
     }
 
-    /**
-     * If there is already a loader task running, tell it to stop.
-     */
-    private void stopLoaderLocked() {
-        LoaderTask oldTask = mLoaderTask;
-        if (oldTask != null) {
-            oldTask.stopLocked();
-        }
-    }
-
     public boolean isCurrentCallbacks(Callbacks callbacks) {
         return (mCallbacks != null && mCallbacks.get() == callbacks);
     }
@@ -484,12 +473,10 @@
                         });
 
                 // If there is already one running, tell it to stop.
-                stopLoaderLocked();
+                stopLoader();
                 LoaderResults loaderResults = new LoaderResults(mApp, sBgDataModel,
                         mBgAllAppsList, synchronousBindPage, mCallbacks);
-                if (synchronousBindPage != PagedView.INVALID_RESTORE_PAGE
-                        && mModelLoaded && !mIsLoaderTaskRunning) {
-
+                if (mModelLoaded && !mIsLoaderTaskRunning) {
                     // Divide the set of loaded items into those that we are binding synchronously,
                     // and everything else that is to be bound normally (asynchronously).
                     loaderResults.bindWorkspace();
@@ -500,22 +487,34 @@
                     loaderResults.bindWidgets();
                     return true;
                 } else {
-                    mLoaderTask = new LoaderTask(mApp, mBgAllAppsList, sBgDataModel, loaderResults);
-                    sWorker.post(mLoaderTask);
+                    startLoaderForResults(loaderResults);
                 }
             }
         }
         return false;
     }
 
+    /**
+     * If there is already a loader task running, tell it to stop.
+     */
     public void stopLoader() {
         synchronized (mLock) {
-            if (mLoaderTask != null) {
-                mLoaderTask.stopLocked();
+            LoaderTask oldTask = mLoaderTask;
+            mLoaderTask = null;
+            if (oldTask != null) {
+                oldTask.stopLocked();
             }
         }
     }
 
+    public void startLoaderForResults(LoaderResults results) {
+        synchronized (mLock) {
+            stopLoader();
+            mLoaderTask = new LoaderTask(mApp, mBgAllAppsList, sBgDataModel, results);
+            runOnWorkerThread(mLoaderTask);
+        }
+    }
+
     /**
      * Loads the workspace screen ids in an ordered list.
      */
@@ -529,7 +528,7 @@
     }
 
     public void onInstallSessionCreated(final PackageInstallInfo sessionInfo) {
-        enqueueModelUpdateTask(new ExtendedModelTask() {
+        enqueueModelUpdateTask(new BaseModelUpdateTask() {
             @Override
             public void execute(LauncherAppState app, BgDataModel dataModel, AllAppsList apps) {
                 apps.addPromiseApp(app.getContext(), sessionInfo);
@@ -607,8 +606,8 @@
                 CacheDataUpdatedTask.OP_CACHE_UPDATE, user, updatedPackages));
     }
 
-    public void enqueueModelUpdateTask(BaseModelUpdateTask task) {
-        task.init(this);
+    public void enqueueModelUpdateTask(ModelUpdateTask task) {
+        task.init(mApp, this, sBgDataModel, mBgAllAppsList, mUiExecutor);
         runOnWorkerThread(task);
     }
 
@@ -624,54 +623,14 @@
     /**
      * A runnable which changes/updates the data model of the launcher based on certain events.
      */
-    public static abstract class BaseModelUpdateTask implements Runnable {
-
-        private LauncherModel mModel;
-        private Executor mUiExecutor;
-
-        /* package private */
-        void init(LauncherModel model) {
-            mModel = model;
-            mUiExecutor = mModel.mUiExecutor;
-        }
-
-        @Override
-        public final void run() {
-            if (!mModel.mModelLoaded) {
-                if (DEBUG_TASKS) {
-                    Log.d(TAG, "Ignoring model task since loader is pending=" + this);
-                }
-                // Loader has not yet run.
-                return;
-            }
-            execute(mModel.mApp, sBgDataModel, mModel.mBgAllAppsList);
-        }
+    public interface ModelUpdateTask extends Runnable {
 
         /**
-         * Execute the actual task. Called on the worker thread.
+         * Called before the task is posted to initialize the internal state.
          */
-        public abstract void execute(
-                LauncherAppState app, BgDataModel dataModel, AllAppsList apps);
+        void init(LauncherAppState app, LauncherModel model,
+                BgDataModel dataModel, AllAppsList allAppsList, Executor uiExecutor);
 
-        /**
-         * Schedules a {@param task} to be executed on the current callbacks.
-         */
-        public final void scheduleCallbackTask(final CallbackTask task) {
-            final Callbacks callbacks = mModel.getCallback();
-            mUiExecutor.execute(new Runnable() {
-                public void run() {
-                    Callbacks cb = mModel.getCallback();
-                    if (callbacks == cb && cb != null) {
-                        task.execute(callbacks);
-                    }
-                }
-            });
-        }
-
-        public ModelWriter getModelWriter() {
-            // Updates from model task, do not deal with icon position in hotseat.
-            return mModel.getWriter(false /* hasVerticalHotseat */);
-        }
     }
 
     public void updateAndBindShortcutInfo(final ShortcutInfo si, final ShortcutInfoCompat info) {
@@ -689,7 +648,7 @@
      * Utility method to update a shortcut on the background thread.
      */
     public void updateAndBindShortcutInfo(final Provider<ShortcutInfo> shortcutProvider) {
-        enqueueModelUpdateTask(new ExtendedModelTask() {
+        enqueueModelUpdateTask(new BaseModelUpdateTask() {
             @Override
             public void execute(LauncherAppState app, BgDataModel dataModel, AllAppsList apps) {
                 ShortcutInfo info = shortcutProvider.get();
@@ -701,7 +660,7 @@
     }
 
     public void refreshAndBindWidgetsAndShortcuts(@Nullable final PackageUserKey packageUser) {
-        enqueueModelUpdateTask(new ExtendedModelTask() {
+        enqueueModelUpdateTask(new BaseModelUpdateTask() {
             @Override
             public void execute(LauncherAppState app, BgDataModel dataModel, AllAppsList apps) {
                 dataModel.widgetsModel.update(app, packageUser);
diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index 2e8e15b..b27ccfd 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -47,7 +47,7 @@
 /**
  * Task to add auto-created workspace items.
  */
-public class AddWorkspaceItemsTask extends ExtendedModelTask {
+public class AddWorkspaceItemsTask extends BaseModelUpdateTask {
 
     private final Provider<List<Pair<ItemInfo, Object>>> mAppsProvider;
 
diff --git a/src/com/android/launcher3/model/BaseModelUpdateTask.java b/src/com/android/launcher3/model/BaseModelUpdateTask.java
new file mode 100644
index 0000000..9b4510f
--- /dev/null
+++ b/src/com/android/launcher3/model/BaseModelUpdateTask.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2016 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.model;
+
+import android.os.UserHandle;
+import android.util.Log;
+
+import com.android.launcher3.AllAppsList;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
+import com.android.launcher3.LauncherModel.CallbackTask;
+import com.android.launcher3.LauncherModel.Callbacks;
+import com.android.launcher3.ShortcutInfo;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.MultiHashMap;
+
+import java.util.ArrayList;
+import java.util.concurrent.Executor;
+
+/**
+ * Extension of {@link ModelUpdateTask} with some utility methods
+ */
+public abstract class BaseModelUpdateTask implements ModelUpdateTask {
+
+    private static final boolean DEBUG_TASKS = false;
+    private static final String TAG = "BaseModelUpdateTask";
+
+    private LauncherAppState mApp;
+    private LauncherModel mModel;
+    private BgDataModel mDataModel;
+    private AllAppsList mAllAppsList;
+    private Executor mUiExecutor;
+
+    public void init(LauncherAppState app, LauncherModel model,
+            BgDataModel dataModel, AllAppsList allAppsList, Executor uiExecutor) {
+        mApp = app;
+        mModel = model;
+        mDataModel = dataModel;
+        mAllAppsList = allAppsList;
+        mUiExecutor = uiExecutor;
+    }
+
+    @Override
+    public final void run() {
+        if (!mModel.isModelLoaded()) {
+            if (DEBUG_TASKS) {
+                Log.d(TAG, "Ignoring model task since loader is pending=" + this);
+            }
+            // Loader has not yet run.
+            return;
+        }
+        execute(mApp, mDataModel, mAllAppsList);
+    }
+
+    /**
+     * Execute the actual task. Called on the worker thread.
+     */
+    public abstract void execute(
+            LauncherAppState app, BgDataModel dataModel, AllAppsList apps);
+
+    /**
+     * Schedules a {@param task} to be executed on the current callbacks.
+     */
+    public final void scheduleCallbackTask(final CallbackTask task) {
+        final Callbacks callbacks = mModel.getCallback();
+        mUiExecutor.execute(new Runnable() {
+            public void run() {
+                Callbacks cb = mModel.getCallback();
+                if (callbacks == cb && cb != null) {
+                    task.execute(callbacks);
+                }
+            }
+        });
+    }
+
+    public ModelWriter getModelWriter() {
+        // Updates from model task, do not deal with icon position in hotseat.
+        return mModel.getWriter(false /* hasVerticalHotseat */);
+    }
+
+
+    public void bindUpdatedShortcuts(
+            ArrayList<ShortcutInfo> updatedShortcuts, UserHandle user) {
+        bindUpdatedShortcuts(updatedShortcuts, new ArrayList<ShortcutInfo>(), user);
+    }
+
+    public void bindUpdatedShortcuts(
+            final ArrayList<ShortcutInfo> updatedShortcuts,
+            final ArrayList<ShortcutInfo> removedShortcuts,
+            final UserHandle user) {
+        if (!updatedShortcuts.isEmpty() || !removedShortcuts.isEmpty()) {
+            scheduleCallbackTask(new CallbackTask() {
+                @Override
+                public void execute(Callbacks callbacks) {
+                    callbacks.bindShortcutsChanged(updatedShortcuts, removedShortcuts, user);
+                }
+            });
+        }
+    }
+
+    public void bindDeepShortcuts(BgDataModel dataModel) {
+        final MultiHashMap<ComponentKey, String> shortcutMapCopy = dataModel.deepShortcutMap.clone();
+        scheduleCallbackTask(new CallbackTask() {
+            @Override
+            public void execute(Callbacks callbacks) {
+                callbacks.bindDeepShortcutMap(shortcutMapCopy);
+            }
+        });
+    }
+
+    public void bindUpdatedWidgets(BgDataModel dataModel) {
+        final MultiHashMap<PackageItemInfo, WidgetItem> widgets
+                = dataModel.widgetsModel.getWidgetsMap();
+        scheduleCallbackTask(new CallbackTask() {
+            @Override
+            public void execute(Callbacks callbacks) {
+                callbacks.bindAllWidgets(widgets);
+            }
+        });
+    }
+}
diff --git a/src/com/android/launcher3/model/CacheDataUpdatedTask.java b/src/com/android/launcher3/model/CacheDataUpdatedTask.java
index d7cece4..8597e10 100644
--- a/src/com/android/launcher3/model/CacheDataUpdatedTask.java
+++ b/src/com/android/launcher3/model/CacheDataUpdatedTask.java
@@ -34,7 +34,7 @@
 /**
  * Handles changes due to cache updates.
  */
-public class CacheDataUpdatedTask extends ExtendedModelTask {
+public class CacheDataUpdatedTask extends BaseModelUpdateTask {
 
     public static final int OP_CACHE_UPDATE = 1;
     public static final int OP_SESSION_UPDATE = 2;
diff --git a/src/com/android/launcher3/model/ExtendedModelTask.java b/src/com/android/launcher3/model/ExtendedModelTask.java
deleted file mode 100644
index 080aaf5..0000000
--- a/src/com/android/launcher3/model/ExtendedModelTask.java
+++ /dev/null
@@ -1,73 +0,0 @@
-/*
- * Copyright (C) 2016 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.model;
-
-import android.os.UserHandle;
-
-import com.android.launcher3.LauncherModel.BaseModelUpdateTask;
-import com.android.launcher3.LauncherModel.CallbackTask;
-import com.android.launcher3.LauncherModel.Callbacks;
-import com.android.launcher3.ShortcutInfo;
-import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.MultiHashMap;
-
-import java.util.ArrayList;
-
-/**
- * Extension of {@link BaseModelUpdateTask} with some utility methods
- */
-public abstract class ExtendedModelTask extends BaseModelUpdateTask {
-
-    public void bindUpdatedShortcuts(
-            ArrayList<ShortcutInfo> updatedShortcuts, UserHandle user) {
-        bindUpdatedShortcuts(updatedShortcuts, new ArrayList<ShortcutInfo>(), user);
-    }
-
-    public void bindUpdatedShortcuts(
-            final ArrayList<ShortcutInfo> updatedShortcuts,
-            final ArrayList<ShortcutInfo> removedShortcuts,
-            final UserHandle user) {
-        if (!updatedShortcuts.isEmpty() || !removedShortcuts.isEmpty()) {
-            scheduleCallbackTask(new CallbackTask() {
-                @Override
-                public void execute(Callbacks callbacks) {
-                    callbacks.bindShortcutsChanged(updatedShortcuts, removedShortcuts, user);
-                }
-            });
-        }
-    }
-
-    public void bindDeepShortcuts(BgDataModel dataModel) {
-        final MultiHashMap<ComponentKey, String> shortcutMapCopy = dataModel.deepShortcutMap.clone();
-        scheduleCallbackTask(new CallbackTask() {
-            @Override
-            public void execute(Callbacks callbacks) {
-                callbacks.bindDeepShortcutMap(shortcutMapCopy);
-            }
-        });
-    }
-
-    public void bindUpdatedWidgets(BgDataModel dataModel) {
-        final MultiHashMap<PackageItemInfo, WidgetItem> widgets
-                = dataModel.widgetsModel.getWidgetsMap();
-        scheduleCallbackTask(new CallbackTask() {
-            @Override
-            public void execute(Callbacks callbacks) {
-                callbacks.bindAllWidgets(widgets);
-            }
-        });
-    }
-}
diff --git a/src/com/android/launcher3/model/LoaderResults.java b/src/com/android/launcher3/model/LoaderResults.java
index 28df64d..0df8b6f 100644
--- a/src/com/android/launcher3/model/LoaderResults.java
+++ b/src/com/android/launcher3/model/LoaderResults.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.model;
 
+import android.os.Looper;
 import android.util.Log;
 
 import com.android.launcher3.AllAppsList;
@@ -32,6 +33,7 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.ViewOnDrawExecutor;
 
@@ -45,7 +47,7 @@
 import java.util.concurrent.Executor;
 
 /**
- * Helper class to handle results of {@link com.android.launcher3.LauncherModel.LoaderTask}.
+ * Helper class to handle results of {@link com.android.launcher3.model.LoaderTask}.
  */
 public class LoaderResults {
 
@@ -389,4 +391,13 @@
         };
         mUiExecutor.execute(r);
     }
+
+    public LooperIdleLock newIdleLock(Object lock) {
+        LooperIdleLock idleLock = new LooperIdleLock(lock, Looper.getMainLooper());
+        // If we are not binding, there is no reason to wait for idle.
+        if (mCallbacks.get() == null) {
+            idleLock.queueIdle();
+        }
+        return idleLock;
+    }
 }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index bcf516e..b24d682 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -116,11 +116,11 @@
         mIconCache = mApp.getIconCache();
     }
 
-    private synchronized void waitForIdle() {
+    protected synchronized void waitForIdle() {
         // Wait until the either we're stopped or the other threads are done.
         // This way we don't start loading all apps until the workspace has settled
         // down.
-        LooperIdleLock idleLock = new LooperIdleLock(this, Looper.getMainLooper());
+        LooperIdleLock idleLock = mResults.newIdleLock(this);
         // Just in case mFlushingWorkerThread changes but we aren't woken up,
         // wait no longer than 1sec at a time
         while (!mStopped && idleLock.awaitLocked(1000));
@@ -202,7 +202,10 @@
 
             transaction.commit();
         } catch (CancellationException e) {
-          // Loader stopped, ignore
+            // Loader stopped, ignore
+            if (DEBUG_LOADERS) {
+                Log.d(TAG, "Loader cancelled", e);
+            }
         }
     }
 
diff --git a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
index 76b90a8..1e0af68 100644
--- a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
+++ b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
@@ -35,7 +35,7 @@
 /**
  * Handles changes due to a sessions updates for a currently installing app.
  */
-public class PackageInstallStateChangedTask extends ExtendedModelTask {
+public class PackageInstallStateChangedTask extends BaseModelUpdateTask {
 
     private final PackageInstallInfo mInstallInfo;
 
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index 46fea21..a237b3e 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -30,7 +30,6 @@
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetInfo;
-import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherModel.CallbackTask;
 import com.android.launcher3.LauncherModel.Callbacks;
 import com.android.launcher3.LauncherSettings;
@@ -44,7 +43,6 @@
 import com.android.launcher3.graphics.LauncherIcons;
 import com.android.launcher3.util.FlagOp;
 import com.android.launcher3.util.ItemInfoMatcher;
-import com.android.launcher3.util.MultiHashMap;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
 
@@ -58,7 +56,7 @@
  * Handles updates due to changes in package manager (app installed/updated/removed)
  * or when a user availability changes.
  */
-public class PackageUpdatedTask extends ExtendedModelTask {
+public class PackageUpdatedTask extends BaseModelUpdateTask {
 
     private static final boolean DEBUG = false;
     private static final String TAG = "PackageUpdatedTask";
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.java b/src/com/android/launcher3/model/ShortcutsChangedTask.java
index 47e83e5..6f32585 100644
--- a/src/com/android/launcher3/model/ShortcutsChangedTask.java
+++ b/src/com/android/launcher3/model/ShortcutsChangedTask.java
@@ -34,7 +34,7 @@
 /**
  * Handles changes due to shortcut manager updates (deep shortcut changes)
  */
-public class ShortcutsChangedTask extends ExtendedModelTask {
+public class ShortcutsChangedTask extends BaseModelUpdateTask {
 
     private final String mPackageName;
     private final List<ShortcutInfoCompat> mShortcuts;
diff --git a/src/com/android/launcher3/model/UserLockStateChangedTask.java b/src/com/android/launcher3/model/UserLockStateChangedTask.java
index fefed75..5682006 100644
--- a/src/com/android/launcher3/model/UserLockStateChangedTask.java
+++ b/src/com/android/launcher3/model/UserLockStateChangedTask.java
@@ -38,7 +38,7 @@
 /**
  * Task to handle changing of lock state of the user
  */
-public class UserLockStateChangedTask extends ExtendedModelTask {
+public class UserLockStateChangedTask extends BaseModelUpdateTask {
 
     private final UserHandle mUser;
 
diff --git a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
index 13e0986..3d03507 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.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 @@
         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 @@
     /**
      * 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());