diff --git a/quickstep/src/com/android/launcher3/model/AppEventProducer.java b/quickstep/src/com/android/launcher3/model/AppEventProducer.java
index 7209bd8..8dded8f 100644
--- a/quickstep/src/com/android/launcher3/model/AppEventProducer.java
+++ b/quickstep/src/com/android/launcher3/model/AppEventProducer.java
@@ -191,7 +191,7 @@
     @Nullable
     AppTarget toAppTarget(LauncherAtom.ItemInfo info) {
         int iconInfoType = getIconInfoTypeFromItemInfo(info);
-        UserCache userCache = UserCache.getInstance(mContext);
+        UserCache userCache = UserCache.INSTANCE.get(mContext);
         UserHandle userHandle = userCache.getUserProfiles().stream()
                 .filter(user -> userCache.getUserInfo(user).type == iconInfoType)
                 .findFirst()
diff --git a/quickstep/tests/src/com/android/launcher3/model/AppEventProducerTest.java b/quickstep/tests/src/com/android/launcher3/model/AppEventProducerTest.java
index d6c1447..d4dd580 100644
--- a/quickstep/tests/src/com/android/launcher3/model/AppEventProducerTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/AppEventProducerTest.java
@@ -28,7 +28,6 @@
 
 import android.app.prediction.AppTarget;
 import android.content.ComponentName;
-import android.content.Context;
 import android.content.Intent;
 import android.os.Process;
 import android.os.UserHandle;
@@ -39,13 +38,12 @@
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.util.ActivityContextWrapper;
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 import com.android.launcher3.util.UserIconInfo;
-import com.android.launcher3.util.rule.StaticMockitoRule;
 import com.android.systemui.shared.system.SysUiStatsLog;
 
+import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -65,22 +63,24 @@
     private static final UserIconInfo PRIVATE_ICON_INFO =
             new UserIconInfo(PRIVATE_HANDLE, UserIconInfo.TYPE_PRIVATE);
 
-    private Context mContext;
+    private SandboxContext mContext;
     private AppEventProducer mAppEventProducer;
     @Mock
     private UserCache mUserCache;
 
-    @Rule
-    public final StaticMockitoRule mStaticMockitoRule = new StaticMockitoRule(UserCache.class);
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = new ActivityContextWrapper(getApplicationContext());
-        when(UserCache.getInstance(any(Context.class))).thenReturn(mUserCache);
+        mContext = new SandboxContext(getApplicationContext());
+        mContext.putObject(UserCache.INSTANCE, mUserCache);
         mAppEventProducer = new AppEventProducer(mContext, null);
     }
 
+    @After
+    public void tearDown() {
+        mContext.onDestroy();
+    }
+
     @Test
     public void buildAppTarget_containsCorrectUser() {
         when(mUserCache.getUserProfiles())
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index d4bcd24..e015021 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -44,6 +44,7 @@
 import com.android.launcher3.icons.IconProvider;
 import com.android.launcher3.icons.LauncherIconProvider;
 import com.android.launcher3.icons.LauncherIcons;
+import com.android.launcher3.model.ModelLauncherCallbacks;
 import com.android.launcher3.notification.NotificationListener;
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.InstallSessionTracker;
@@ -101,9 +102,10 @@
             }
         });
 
-        mContext.getSystemService(LauncherApps.class).registerCallback(mModel);
+        ModelLauncherCallbacks callbacks = mModel.newModelCallbacks();
+        mContext.getSystemService(LauncherApps.class).registerCallback(callbacks);
         mOnTerminateCallback.add(() ->
-                mContext.getSystemService(LauncherApps.class).unregisterCallback(mModel));
+                mContext.getSystemService(LauncherApps.class).unregisterCallback(callbacks));
 
         SimpleBroadcastReceiver modelChangeReceiver =
                 new SimpleBroadcastReceiver(mModel::onBroadcastIntent);
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index c81db63..d124746 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -20,6 +20,7 @@
 
 import static com.android.launcher3.LauncherAppState.ACTION_FORCE_ROLOAD;
 import static com.android.launcher3.config.FeatureFlags.IS_STUDIO_BUILD;
+import static com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE;
 import static com.android.launcher3.pm.UserCache.ACTION_PROFILE_AVAILABLE;
 import static com.android.launcher3.pm.UserCache.ACTION_PROFILE_UNAVAILABLE;
 import static com.android.launcher3.testing.shared.TestProtocol.sDebugTracing;
@@ -28,7 +29,6 @@
 
 import android.content.Context;
 import android.content.Intent;
-import android.content.pm.LauncherApps;
 import android.content.pm.PackageInstaller;
 import android.content.pm.ShortcutInfo;
 import android.os.UserHandle;
@@ -43,7 +43,6 @@
 import com.android.launcher3.celllayout.CellPosMapper;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.AddWorkspaceItemsTask;
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BaseModelUpdateTask;
@@ -55,8 +54,8 @@
 import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.model.ModelDelegate;
+import com.android.launcher3.model.ModelLauncherCallbacks;
 import com.android.launcher3.model.ModelWriter;
-import com.android.launcher3.model.PackageIncrementalDownloadUpdatedTask;
 import com.android.launcher3.model.PackageInstallStateChangedTask;
 import com.android.launcher3.model.PackageUpdatedTask;
 import com.android.launcher3.model.ReloadStringCacheTask;
@@ -89,7 +88,7 @@
  * LauncherModel object held in a static. Also provide APIs for updating the database state
  * for the Launcher.
  */
-public class LauncherModel extends LauncherApps.Callback implements InstallSessionTracker.Callback {
+public class LauncherModel implements InstallSessionTracker.Callback {
     private static final boolean DEBUG_RECEIVER = false;
 
     static final String TAG = "Launcher.Model";
@@ -168,6 +167,10 @@
         return mModelDbController;
     }
 
+    public ModelLauncherCallbacks newModelCallbacks() {
+        return new ModelLauncherCallbacks(this::enqueueModelUpdateTask);
+    }
+
     /**
      * Adds the provided items to the workspace.
      */
@@ -186,77 +189,6 @@
                 owner);
     }
 
-    @Override
-    public void onPackageChanged(
-            @NonNull final String packageName, @NonNull final UserHandle user) {
-        int op = PackageUpdatedTask.OP_UPDATE;
-        enqueueModelUpdateTask(new PackageUpdatedTask(op, user, packageName));
-    }
-
-    @Override
-    public void onPackageRemoved(
-            @NonNull final String packageName, @NonNull final UserHandle user) {
-        onPackagesRemoved(user, packageName);
-    }
-
-    public void onPackagesRemoved(
-            @NonNull final UserHandle user, @NonNull final String... packages) {
-        int op = PackageUpdatedTask.OP_REMOVE;
-        FileLog.d(TAG, "package removed received " + TextUtils.join(",", packages));
-        enqueueModelUpdateTask(new PackageUpdatedTask(op, user, packages));
-    }
-
-    @Override
-    public void onPackageAdded(@NonNull final String packageName, @NonNull final UserHandle user) {
-        int op = PackageUpdatedTask.OP_ADD;
-        enqueueModelUpdateTask(new PackageUpdatedTask(op, user, packageName));
-    }
-
-    @Override
-    public void onPackagesAvailable(@NonNull final String[] packageNames,
-            @NonNull final UserHandle user, final boolean replacing) {
-        enqueueModelUpdateTask(
-                new PackageUpdatedTask(PackageUpdatedTask.OP_UPDATE, user, packageNames));
-    }
-
-    @Override
-    public void onPackagesUnavailable(@NonNull final String[] packageNames,
-            @NonNull final UserHandle user, final boolean replacing) {
-        if (!replacing) {
-            enqueueModelUpdateTask(new PackageUpdatedTask(
-                    PackageUpdatedTask.OP_UNAVAILABLE, user, packageNames));
-        }
-    }
-
-    @Override
-    public void onPackagesSuspended(
-            @NonNull final String[] packageNames, @NonNull final UserHandle user) {
-        enqueueModelUpdateTask(new PackageUpdatedTask(
-                PackageUpdatedTask.OP_SUSPEND, user, packageNames));
-    }
-
-    @Override
-    public void onPackagesUnsuspended(
-            @NonNull final String[] packageNames, @NonNull final UserHandle user) {
-        enqueueModelUpdateTask(new PackageUpdatedTask(
-                PackageUpdatedTask.OP_UNSUSPEND, user, packageNames));
-    }
-
-    @Override
-    public void onPackageLoadingProgressChanged(@NonNull final String packageName,
-            @NonNull final UserHandle user, final float progress) {
-        if (Utilities.ATLEAST_S) {
-            enqueueModelUpdateTask(new PackageIncrementalDownloadUpdatedTask(
-                    packageName, user, progress));
-        }
-    }
-
-    @Override
-    public void onShortcutsChanged(@NonNull final String packageName,
-            @NonNull final List<ShortcutInfo> shortcuts, @NonNull final UserHandle user) {
-        enqueueModelUpdateTask(new ShortcutsChangedTask(packageName, shortcuts, user, true));
-    }
-
     /**
      * Called when the icon for an app changes, outside of package event
      */
@@ -265,7 +197,7 @@
             @NonNull final UserHandle user) {
         // Update the icon for the calendar package
         Context context = mApp.getContext();
-        onPackageChanged(packageName, user);
+        enqueueModelUpdateTask(new PackageUpdatedTask(OP_UPDATE, user, packageName));
 
         List<ShortcutInfo> pinnedShortcuts = new ShortcutRequest(context, user)
                 .forPackage(packageName).query(ShortcutRequest.PINNED);
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 7b9f6fa..905d03a 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -123,6 +123,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CancellationException;
 
@@ -185,7 +186,7 @@
         mLauncherBinder = launcherBinder;
         mLauncherApps = mApp.getContext().getSystemService(LauncherApps.class);
         mUserManager = mApp.getContext().getSystemService(UserManager.class);
-        mUserCache = UserCache.getInstance(mApp.getContext());
+        mUserCache = UserCache.INSTANCE.get(mApp.getContext());
         mSessionHelper = InstallSessionHelper.INSTANCE.get(mApp.getContext());
         mIconCache = mApp.getIconCache();
         mUserManagerState = userManagerState;
@@ -247,7 +248,7 @@
             // sanitizeData should not be invoked if the workspace is loaded from a db different
             // from the main db as defined in the invariant device profile.
             // (e.g. both grid preview and minimal device mode uses a different db)
-            if (mApp.getInvariantDeviceProfile().dbFile.equals(mDbName)) {
+            if (Objects.equals(mApp.getInvariantDeviceProfile().dbFile, mDbName)) {
                 verifyNotStopped();
                 sanitizeFolders(mItemsDeleted);
                 sanitizeWidgetsShortcutsAndPackages();
diff --git a/src/com/android/launcher3/model/ModelLauncherCallbacks.kt b/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
new file mode 100644
index 0000000..b12b2bc
--- /dev/null
+++ b/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
@@ -0,0 +1,103 @@
+/*
+ * Copyright (C) 2024 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.content.pm.LauncherApps
+import android.content.pm.ShortcutInfo
+import android.os.UserHandle
+import android.text.TextUtils
+import com.android.launcher3.LauncherModel.ModelUpdateTask
+import com.android.launcher3.logging.FileLog
+import com.android.launcher3.model.PackageUpdatedTask.OP_ADD
+import com.android.launcher3.model.PackageUpdatedTask.OP_REMOVE
+import com.android.launcher3.model.PackageUpdatedTask.OP_SUSPEND
+import com.android.launcher3.model.PackageUpdatedTask.OP_UNAVAILABLE
+import com.android.launcher3.model.PackageUpdatedTask.OP_UNSUSPEND
+import com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE
+import java.util.function.Consumer
+
+/**
+ * Implementation of {@link LauncherApps#Callbacks} which converts various events to corresponding
+ * model tasks
+ */
+class ModelLauncherCallbacks(private var taskExecutor: Consumer<ModelUpdateTask>) :
+    LauncherApps.Callback() {
+
+    override fun onPackageAdded(packageName: String, user: UserHandle) {
+        taskExecutor.accept(PackageUpdatedTask(OP_ADD, user, packageName))
+    }
+
+    override fun onPackageChanged(packageName: String, user: UserHandle) {
+        taskExecutor.accept(PackageUpdatedTask(OP_UPDATE, user, packageName))
+    }
+
+    override fun onPackageLoadingProgressChanged(
+        packageName: String,
+        user: UserHandle,
+        progress: Float
+    ) {
+        taskExecutor.accept(PackageIncrementalDownloadUpdatedTask(packageName, user, progress))
+    }
+
+    override fun onPackageRemoved(packageName: String, user: UserHandle) {
+        FileLog.d(TAG, "package removed received $packageName")
+        taskExecutor.accept(PackageUpdatedTask(OP_REMOVE, user, packageName))
+    }
+
+    override fun onPackagesAvailable(
+        vararg packageNames: String,
+        user: UserHandle,
+        replacing: Boolean
+    ) {
+        taskExecutor.accept(PackageUpdatedTask(OP_UPDATE, user, *packageNames))
+    }
+
+    override fun onPackagesSuspended(vararg packageNames: String, user: UserHandle) {
+        taskExecutor.accept(PackageUpdatedTask(OP_SUSPEND, user, *packageNames))
+    }
+
+    override fun onPackagesUnavailable(
+        packageNames: Array<String>,
+        user: UserHandle,
+        replacing: Boolean
+    ) {
+        if (!replacing) {
+            taskExecutor.accept(PackageUpdatedTask(OP_UNAVAILABLE, user, *packageNames))
+        }
+    }
+
+    override fun onPackagesUnsuspended(vararg packageNames: String, user: UserHandle) {
+        taskExecutor.accept(PackageUpdatedTask(OP_UNSUSPEND, user, *packageNames))
+    }
+
+    override fun onShortcutsChanged(
+        packageName: String,
+        shortcuts: MutableList<ShortcutInfo>,
+        user: UserHandle
+    ) {
+        taskExecutor.accept(ShortcutsChangedTask(packageName, shortcuts, user, true))
+    }
+
+    fun onPackagesRemoved(user: UserHandle, packages: List<String>) {
+        FileLog.d(TAG, "package removed received " + TextUtils.join(",", packages))
+        taskExecutor.accept(PackageUpdatedTask(OP_REMOVE, user, *packages.toTypedArray()))
+    }
+
+    companion object {
+        private const val TAG = "LauncherAppsCallbackImpl"
+    }
+}
diff --git a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
index 76a87ed..2457a42 100644
--- a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
+++ b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
@@ -52,7 +52,8 @@
                 ApplicationInfo ai = app.getContext()
                         .getPackageManager().getApplicationInfo(mInstallInfo.packageName, 0);
                 if (InstantAppResolver.newInstance(app.getContext()).isInstantApp(ai)) {
-                    app.getModel().onPackageAdded(ai.packageName, mInstallInfo.user);
+                    app.getModel().newModelCallbacks()
+                            .onPackageAdded(ai.packageName, mInstallInfo.user);
                 }
             } catch (PackageManager.NameNotFoundException e) {
                 // Ignore
diff --git a/src/com/android/launcher3/model/SdCardAvailableReceiver.java b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
index 3798575..8cfa3aa 100644
--- a/src/com/android/launcher3/model/SdCardAvailableReceiver.java
+++ b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
@@ -67,11 +67,10 @@
                 }
             }
             if (!packagesRemoved.isEmpty()) {
-                mModel.onPackagesRemoved(user,
-                        packagesRemoved.toArray(new String[packagesRemoved.size()]));
+                mModel.newModelCallbacks().onPackagesRemoved(user, packagesRemoved);
             }
             if (!packagesUnavailable.isEmpty()) {
-                mModel.onPackagesUnavailable(
+                mModel.newModelCallbacks().onPackagesUnavailable(
                         packagesUnavailable.toArray(new String[packagesUnavailable.size()]),
                         user, false);
             }
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index def27b8..dbca9d1 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -1,6 +1,5 @@
 package com.android.launcher3.model
 
-import android.content.Context
 import android.os.UserHandle
 import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
@@ -11,16 +10,20 @@
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherModel
 import com.android.launcher3.LauncherModel.LoaderTransaction
+import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.icons.IconCache
 import com.android.launcher3.icons.cache.CachingLogic
 import com.android.launcher3.icons.cache.IconCacheUpdateHandler
+import com.android.launcher3.pm.InstallSessionHelper
 import com.android.launcher3.pm.UserCache
+import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.LooperIdleLock
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.android.launcher3.util.UserIconInfo
-import com.android.launcher3.util.rule.StaticMockitoRule
 import com.google.common.truth.Truth
 import java.util.concurrent.CountDownLatch
+import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -39,6 +42,7 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class LoaderTaskTest {
+    private lateinit var context: SandboxContext
     @Mock private lateinit var app: LauncherAppState
     @Mock private lateinit var bgAllAppsList: AllAppsList
     @Mock private lateinit var modelDelegate: ModelDelegate
@@ -52,21 +56,28 @@
 
     @Spy private var userManagerState: UserManagerState? = UserManagerState()
 
-    @get:Rule(order = 0) val staticMockitoRule = StaticMockitoRule(UserCache::class.java)
-    @get:Rule(order = 1)
-    val setFlagsRule = SetFlagsRule().apply { initAllFlagsToReleaseConfigDefault() }
+    @get:Rule val setFlagsRule = SetFlagsRule().apply { initAllFlagsToReleaseConfigDefault() }
 
     @Before
     fun setup() {
-        val context = InstrumentationRegistry.getInstrumentation().targetContext
+        MockitoAnnotations.initMocks(this)
+
+        context =
+            SandboxContext(
+                InstrumentationRegistry.getInstrumentation().targetContext,
+                InstallSessionHelper.INSTANCE,
+                LauncherPrefs.INSTANCE,
+                ItemInstallQueue.INSTANCE,
+                PluginManagerWrapper.INSTANCE
+            )
         val idp =
-            InvariantDeviceProfile.INSTANCE[context].apply {
+            InvariantDeviceProfile().apply {
                 numRows = 5
                 numColumns = 6
                 numDatabaseHotseatIcons = 5
             }
+        context.putObject(InvariantDeviceProfile.INSTANCE, idp)
 
-        MockitoAnnotations.initMocks(this)
         `when`(app.context).thenReturn(context)
         `when`(app.model).thenReturn(launcherModel)
         `when`(launcherModel.beginLoader(any(LoaderTask::class.java))).thenReturn(transaction)
@@ -77,7 +88,12 @@
         `when`(launcherBinder.newIdleLock(any(LoaderTask::class.java))).thenReturn(idleLock)
         `when`(idleLock.awaitLocked(1000)).thenReturn(false)
         `when`(iconCache.updateHandler).thenReturn(iconCacheUpdateHandler)
-        `when`(UserCache.getInstance(any(Context::class.java))).thenReturn(userCache)
+        context.putObject(UserCache.INSTANCE, userCache)
+    }
+
+    @After
+    fun tearDown() {
+        context.onDestroy()
     }
 
     @Test
diff --git a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
index 3dfd6b4..25a4c4e 100644
--- a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
+++ b/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -73,6 +73,10 @@
         TestUtil.uninstallDummyApp();
     }
 
+    private ModelLauncherCallbacks getCallbacks() {
+        return mModelHelper.getModel().newModelCallbacks();
+    }
+
     @Test
     public void testTwoCallbacks_loadedTogether() throws Exception {
         setupWorkspacePages(3);
@@ -127,14 +131,14 @@
 
         // Install package 1
         TestUtil.installDummyApp();
-        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        getCallbacks().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
         assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
         assertTrue(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
 
         // Uninstall package 2
         TestUtil.uninstallDummyApp();
-        mModelHelper.getModel().onPackageRemoved(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        getCallbacks().onPackageRemoved(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
         assertFalse(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
         assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
@@ -142,7 +146,7 @@
         // Unregister a callback and verify updates no longer received
         Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
         TestUtil.installDummyApp();
-        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        getCallbacks().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
 
         // cb2 didn't get the update
diff --git a/tests/src/com/android/launcher3/util/rule/StaticMockitoRule.java b/tests/src/com/android/launcher3/util/rule/StaticMockitoRule.java
deleted file mode 100644
index 6b91474..0000000
--- a/tests/src/com/android/launcher3/util/rule/StaticMockitoRule.java
+++ /dev/null
@@ -1,72 +0,0 @@
-/*
- * Copyright (C) 2022 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 static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
-
-import com.android.dx.mockito.inline.extended.StaticMockitoSession;
-import com.android.dx.mockito.inline.extended.StaticMockitoSessionBuilder;
-
-import org.junit.rules.MethodRule;
-import org.junit.runners.model.FrameworkMethod;
-import org.junit.runners.model.Statement;
-import org.mockito.junit.MockitoRule;
-import org.mockito.quality.Strictness;
-
-/**
- * Similar to {@link MockitoRule}, but uses {@link StaticMockitoSession}, which allows mocking
- * static methods.
- */
-public class StaticMockitoRule implements MethodRule {
-    private Class<?>[] mClasses;
-
-    public StaticMockitoRule(Class<?>... classes) {
-        mClasses = classes;
-    }
-
-    @Override
-    public Statement apply(Statement base, FrameworkMethod method, Object target) {
-        return new Statement() {
-            public void evaluate() throws Throwable {
-                StaticMockitoSessionBuilder builder =
-                        mockitoSession()
-                                .name(target.getClass().getSimpleName() + "." + method.getName())
-                                .initMocks(target)
-                                .strictness(Strictness.STRICT_STUBS);
-
-                for (Class<?> clazz : mClasses) {
-                    builder.mockStatic(clazz);
-                }
-
-                StaticMockitoSession session = builder.startMocking();
-                Throwable testFailure = evaluateSafely(base);
-                session.finishMocking(testFailure);
-                if (testFailure != null) {
-                    throw testFailure;
-                }
-            }
-
-            private Throwable evaluateSafely(Statement base) {
-                try {
-                    base.evaluate();
-                    return null;
-                } catch (Throwable throwable) {
-                    return throwable;
-                }
-            }
-        };
-    }
-}