Making ModelDbController injectable

Bug: 361850561
Test: Updated tests
Flag: EXEMPT dagger
Change-Id: Ie600469476ce5eb60813a3c8b57b9445fcf82c24

17 files changed, 220 insertions, 230 deletions

diff --git a/src/com/android/launcher3/LauncherModel.kt b/src/com/android/launcher3/LauncherModel.kt
index 02d70ae386..557ad67386 100644
--- a/src/com/android/launcher3/LauncherModel.kt
+++ b/src/com/android/launcher3/LauncherModel.kt
@@ -86,12 +86,11 @@ constructor(
     private val loaderFactory: LoaderTaskFactory,
     private val binderFactory: BaseLauncherBinderFactory,
     private val spaceFinderFactory: Provider<WorkspaceItemSpaceFinder>,
+    val modelDbController: ModelDbController,
 ) {
 
     private val mCallbacksList = ArrayList<BgDataModel.Callbacks>(1)
 
-    val modelDbController = ModelDbController(context)
-
     private val mLock = Any()
 
     private var mLoaderTask: LoaderTask? = null
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 03ecf14ef6..acb2b48aeb 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -26,13 +26,13 @@ import android.content.ContentUris;
 import android.content.ContentValues;
 import android.content.pm.PackageManager;
 import android.database.Cursor;
-import android.database.sqlite.SQLiteQueryBuilder;
 import android.net.Uri;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.Process;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.Pair;
 
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.ModelDbController;
@@ -74,24 +74,20 @@ public class LauncherProvider extends ContentProvider {
 
     @Override
     public String getType(Uri uri) {
-        SqlArguments args = new SqlArguments(uri, null, null);
-        if (TextUtils.isEmpty(args.where)) {
-            return "vnd.android.cursor.dir/" + args.table;
+        if (TextUtils.isEmpty(parseUri(uri, null, null).first)) {
+            return "vnd.android.cursor.dir/" + Favorites.TABLE_NAME;
         } else {
-            return "vnd.android.cursor.item/" + args.table;
+            return "vnd.android.cursor.item/" + Favorites.TABLE_NAME;
         }
     }
 
     @Override
     public Cursor query(Uri uri, String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
-        SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-        SQLiteQueryBuilder qb = new SQLiteQueryBuilder();
-        qb.setTables(args.table);
-
+        Pair<String, String[]> args = parseUri(uri, selection, selectionArgs);
         Cursor[] result = new Cursor[1];
         executeControllerTask(controller -> {
-            result[0] = controller.query(args.table, projection, args.where, args.args, sortOrder);
+            result[0] = controller.query(projection, args.first, args.second, sortOrder);
             return 0;
         });
         return result[0];
@@ -108,7 +104,7 @@ public class LauncherProvider extends ContentProvider {
             // attempt allocate and bind the widget.
             Integer itemType = values.getAsInteger(Favorites.ITEM_TYPE);
             if (itemType != null
-                    && itemType.intValue() == Favorites.ITEM_TYPE_APPWIDGET
+                    && itemType == Favorites.ITEM_TYPE_APPWIDGET
                     && !values.containsKey(Favorites.APPWIDGET_ID)) {
 
                 ComponentName cn = ComponentName.unflattenFromString(
@@ -135,8 +131,7 @@ public class LauncherProvider extends ContentProvider {
                 }
             }
 
-            SqlArguments args = new SqlArguments(uri);
-            return controller.insert(args.table, values);
+            return controller.insert(values);
         });
 
         return rowId < 0 ? null : ContentUris.withAppendedId(uri, rowId);
@@ -144,14 +139,14 @@ public class LauncherProvider extends ContentProvider {
 
     @Override
     public int delete(Uri uri, String selection, String[] selectionArgs) {
-        SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-        return executeControllerTask(c -> c.delete(args.table, args.where, args.args));
+        Pair<String, String[]> args = parseUri(uri, selection, selectionArgs);
+        return executeControllerTask(c -> c.delete(args.first, args.second));
     }
 
     @Override
     public int update(Uri uri, ContentValues values, String selection, String[] selectionArgs) {
-        SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-        return executeControllerTask(c -> c.update(args.table, values, args.where, args.args));
+        Pair<String, String[]> args = parseUri(uri, selection, selectionArgs);
+        return executeControllerTask(c -> c.update(values, args.first, args.second));
     }
 
     @Override
@@ -209,35 +204,24 @@ public class LauncherProvider extends ContentProvider {
         }
     }
 
-    static class SqlArguments {
-        public final String table;
-        public final String where;
-        public final String[] args;
-
-        SqlArguments(Uri url, String where, String[] args) {
-            if (url.getPathSegments().size() == 1) {
-                this.table = url.getPathSegments().get(0);
-                this.where = where;
-                this.args = args;
-            } else if (url.getPathSegments().size() != 2) {
-                throw new IllegalArgumentException("Invalid URI: " + url);
-            } else if (!TextUtils.isEmpty(where)) {
-                throw new UnsupportedOperationException("WHERE clause not supported: " + url);
-            } else {
-                this.table = url.getPathSegments().get(0);
-                this.where = "_id=" + ContentUris.parseId(url);
-                this.args = null;
+    /**
+     * Parses the uri and returns the where and arg clause.
+     *
+     * Note: This should be called on the binder thread (before posting on any executor) so that
+     * any parsing error gets propagated to the caller.
+     */
+    private static Pair<String, String[]> parseUri(Uri url, String where, String[] args) {
+        switch (url.getPathSegments().size()) {
+            case 1 -> {
+                return Pair.create(where, args);
             }
-        }
-
-        SqlArguments(Uri url) {
-            if (url.getPathSegments().size() == 1) {
-                table = url.getPathSegments().get(0);
-                where = null;
-                args = null;
-            } else {
-                throw new IllegalArgumentException("Invalid URI: " + url);
+            case 2 -> {
+                if (!TextUtils.isEmpty(where)) {
+                    throw new UnsupportedOperationException("WHERE clause not supported: " + url);
+                }
+                return Pair.create("_id=" + ContentUris.parseId(url), null);
             }
+            default -> throw new IllegalArgumentException("Invalid URI: " + url);
         }
     }
 }
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 5a9b9c20f9..457d12e8ca 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -21,7 +21,6 @@ import static android.content.res.Configuration.UI_MODE_NIGHT_YES;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.launcher3.LauncherPrefs.GRID_NAME;
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
@@ -179,7 +178,7 @@ public class PreviewSurfaceRenderer {
 
         ModelDbController mainController =
                 LauncherAppState.getInstance(mContext).getModel().getModelDbController();
-        try (Cursor c = mainController.query(TABLE_NAME,
+        try (Cursor c = mainController.query(
                 new String[] {
                         LauncherSettings.Favorites.APPWIDGET_ID,
                         LauncherSettings.Favorites.SPANX,
diff --git a/src/com/android/launcher3/model/LoaderCursor.java b/src/com/android/launcher3/model/LoaderCursor.java
index 77d0d7656b..efe61572fe 100644
--- a/src/com/android/launcher3/model/LoaderCursor.java
+++ b/src/com/android/launcher3/model/LoaderCursor.java
@@ -21,7 +21,6 @@ import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED;
@@ -466,7 +465,7 @@ public class LoaderCursor extends CursorWrapper {
     public boolean commitDeleted() {
         if (mItemsToRemove.size() > 0) {
             // Remove dead items
-            mModel.getModelDbController().delete(TABLE_NAME,
+            mModel.getModelDbController().delete(
                     Utilities.createDbSelectionQuery(Favorites._ID, mItemsToRemove), null);
             return true;
         }
@@ -492,7 +491,7 @@ public class LoaderCursor extends CursorWrapper {
             // Update restored items that no longer require special handling
             ContentValues values = new ContentValues();
             values.put(Favorites.RESTORED, 0);
-            mModel.getModelDbController().update(TABLE_NAME, values,
+            mModel.getModelDbController().update(values,
                     Utilities.createDbSelectionQuery(Favorites._ID, mRestoredRows), null);
         }
         if (mRestoreEventLogger != null) {
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 73af6a221a..78e5d898cd 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -23,7 +23,6 @@ import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
 import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 import static com.android.launcher3.LauncherPrefs.SHOULD_SHOW_SMARTSPACE;
 import static com.android.launcher3.LauncherSettings.Favorites.DESKTOP_ICON_FLAG;
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.icons.CacheableShortcutInfo.convertShortcutsToCacheableShortcuts;
 import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_HAS_SHORTCUT_PERMISSION;
@@ -485,7 +484,7 @@ public class LoaderTask implements Runnable {
 
             mShortcutKeyToPinnedShortcuts = new HashMap<>();
             final LoaderCursor c = mLoaderCursorFactory.createLoaderCursor(
-                    dbController.query(TABLE_NAME, null, selection, null, null),
+                    dbController.query(null, selection, null, null),
                     mUserManagerState,
                     mIsRestoreFromBackup ? restoreEventLogger : null);
             final Bundle extras = c.getExtras();
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index feae632bc5..64b9c1c9cd 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -73,6 +73,8 @@ import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.backuprestore.LauncherRestoreEventLogger;
 import com.android.launcher3.backuprestore.LauncherRestoreEventLogger.RestoreError;
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.LauncherDbUtils;
@@ -91,10 +93,13 @@ import java.io.StringReader;
 import java.util.List;
 import java.util.stream.Collectors;
 
+import javax.inject.Inject;
+
 /**
  * Utility class which maintains an instance of Launcher database and provides utility methods
  * around it.
  */
+@LauncherAppSingleton
 public class ModelDbController {
     private static final String TAG = "ModelDbController";
 
@@ -105,17 +110,25 @@ public class ModelDbController {
     protected DatabaseHelper mOpenHelper;
 
     private final Context mContext;
-
-    public ModelDbController(Context context) {
+    private final InvariantDeviceProfile mIdp;
+    private final LauncherPrefs mPrefs;
+    private final UserCache mUserCache;
+
+    @Inject
+    ModelDbController(
+            @ApplicationContext Context context,
+            InvariantDeviceProfile idp,
+            LauncherPrefs prefs,
+            UserCache userCache) {
         mContext = context;
+        mIdp = idp;
+        mPrefs = prefs;
+        mUserCache = userCache;
     }
 
     private void printDBs(String prefix) {
         try {
-            File directory = new File(
-                    mContext.getDatabasePath(InvariantDeviceProfile.INSTANCE.get(mContext).dbFile)
-                            .getParent()
-            );
+            File directory = new File(mContext.getDatabasePath(mIdp.dbFile).getParent());
             if (directory.exists()) {
                 for (File file : directory.listFiles()) {
                     Log.d("b/353505773", prefix + "Database file: " + file.getName());
@@ -130,9 +143,9 @@ public class ModelDbController {
 
     private synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
-            String dbFile = LauncherPrefs.get(mContext).get(DB_FILE);
+            String dbFile = mPrefs.get(DB_FILE);
             if (dbFile.isEmpty()) {
-                dbFile = InvariantDeviceProfile.INSTANCE.get(mContext).dbFile;
+                dbFile = mIdp.dbFile;
             }
             mOpenHelper = createDatabaseHelper(false /* forMigration */, dbFile);
             printDBs("before: ");
@@ -144,7 +157,7 @@ public class ModelDbController {
     protected DatabaseHelper createDatabaseHelper(boolean forMigration, String dbFile) {
         // Set the flag for empty DB
         Runnable onEmptyDbCreateCallback = forMigration ? () -> { }
-                : () -> LauncherPrefs.get(mContext).putSync(getEmptyDbCreatedKey(dbFile).to(true));
+                : () -> mPrefs.putSync(getEmptyDbCreatedKey(dbFile).to(true));
 
         DatabaseHelper databaseHelper = new DatabaseHelper(mContext, dbFile,
                 this::getSerialNumberForUser, onEmptyDbCreateCallback);
@@ -169,12 +182,12 @@ public class ModelDbController {
      * Refer {@link SQLiteDatabase#query}
      */
     @WorkerThread
-    public Cursor query(String table, String[] projection, String selection,
+    public Cursor query(String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
         createDbIfNotExists();
         SQLiteDatabase db = mOpenHelper.getWritableDatabase();
         Cursor result = db.query(
-                table, projection, selection, selectionArgs, null, null, sortOrder);
+                TABLE_NAME, projection, selection, selectionArgs, null, null, sortOrder);
 
         final Bundle extra = new Bundle();
         extra.putString(EXTRA_DB_NAME, mOpenHelper.getDatabaseName());
@@ -186,12 +199,12 @@ public class ModelDbController {
      * Refer {@link SQLiteDatabase#insert(String, String, ContentValues)}
      */
     @WorkerThread
-    public int insert(String table, ContentValues initialValues) {
+    public int insert(ContentValues initialValues) {
         createDbIfNotExists();
 
         SQLiteDatabase db = mOpenHelper.getWritableDatabase();
         addModifiedTime(initialValues);
-        int rowId = mOpenHelper.dbInsertAndCheck(db, table, initialValues);
+        int rowId = mOpenHelper.dbInsertAndCheck(db, TABLE_NAME, initialValues);
         if (rowId >= 0) {
             onAddOrDeleteOp(db);
         }
@@ -202,11 +215,11 @@ public class ModelDbController {
      * Refer {@link SQLiteDatabase#delete(String, String, String[])}
      */
     @WorkerThread
-    public int delete(String table, String selection, String[] selectionArgs) {
+    public int delete(String selection, String[] selectionArgs) {
         createDbIfNotExists();
         SQLiteDatabase db = mOpenHelper.getWritableDatabase();
 
-        int count = db.delete(table, selection, selectionArgs);
+        int count = db.delete(TABLE_NAME, selection, selectionArgs);
         if (count > 0) {
             onAddOrDeleteOp(db);
         }
@@ -217,14 +230,12 @@ public class ModelDbController {
      * Refer {@link SQLiteDatabase#update(String, ContentValues, String, String[])}
      */
     @WorkerThread
-    public int update(String table, ContentValues values,
-            String selection, String[] selectionArgs) {
+    public int update(ContentValues values, String selection, String[] selectionArgs) {
         createDbIfNotExists();
 
         addModifiedTime(values);
         SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        int count = db.update(table, values, selection, selectionArgs);
-        return count;
+        return db.update(TABLE_NAME, values, selection, selectionArgs);
     }
 
     /**
@@ -261,7 +272,7 @@ public class ModelDbController {
     public void createEmptyDB() {
         createDbIfNotExists();
         mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
-        LauncherPrefs.get(mContext).putSync(getEmptyDbCreatedKey().to(true));
+        mPrefs.putSync(getEmptyDbCreatedKey().to(true));
     }
 
     /**
@@ -292,7 +303,6 @@ public class ModelDbController {
                 mOpenHelper.getReadableDatabase(), Favorites.HYBRID_HOTSEAT_BACKUP_TABLE);
     }
 
-
     /**
      * Resets the launcher DB if we should reset it.
      */
@@ -302,11 +312,10 @@ public class ModelDbController {
         }
         FileLog.d(TAG, "resetLauncherDb: Migration failed: resetting launcher database");
         createEmptyDB();
-        LauncherPrefs.get(mContext).putSync(
-                getEmptyDbCreatedKey(mOpenHelper.getDatabaseName()).to(true));
+        mPrefs.putSync(getEmptyDbCreatedKey(mOpenHelper.getDatabaseName()).to(true));
 
         // Write the grid state to avoid another migration
-        new DeviceGridState(LauncherAppState.getIDP(mContext)).writeToPrefs(mContext);
+        new DeviceGridState(mIdp).writeToPrefs(mContext);
     }
 
     /**
@@ -326,7 +335,7 @@ public class ModelDbController {
     }
 
     private boolean isThereExistingDb() {
-        if (LauncherPrefs.get(mContext).get(getEmptyDbCreatedKey())) {
+        if (mPrefs.get(getEmptyDbCreatedKey())) {
             // If we already have a new DB, ignore migration
             FileLog.d(TAG, "isThereExistingDb: new DB already created, skipping migration");
             return true;
@@ -335,8 +344,7 @@ public class ModelDbController {
     }
 
     private boolean isGridMigrationNecessary() {
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
-        if (GridSizeMigrationDBController.needsToMigrate(mContext, idp)) {
+        if (GridSizeMigrationDBController.needsToMigrate(mContext, mIdp)) {
             return true;
         }
         FileLog.d(TAG, "isGridMigrationNecessary: no grid migration needed");
@@ -344,8 +352,7 @@ public class ModelDbController {
     }
 
     private boolean isCurrentDbSameAsTarget() {
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
-        String targetDbName = new DeviceGridState(idp).getDbFile();
+        String targetDbName = new DeviceGridState(mIdp).getDbFile();
         if (TextUtils.equals(targetDbName, mOpenHelper.getDatabaseName())) {
             FileLog.e(TAG, "isCurrentDbSameAsTarget: target db is same as current"
                     + " current db: " + mOpenHelper.getDatabaseName()
@@ -367,7 +374,6 @@ public class ModelDbController {
             return;
         }
 
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
         DatabaseHelper oldHelper = mOpenHelper;
 
         // We save the existing db's before creating the destination db helper so we know what logic
@@ -376,12 +382,12 @@ public class ModelDbController {
                 .filter(dbName -> mContext.getDatabasePath(dbName).exists())
                 .collect(Collectors.toList());
 
-        mOpenHelper = createDatabaseHelper(true, new DeviceGridState(idp).getDbFile());
+        mOpenHelper = createDatabaseHelper(true, new DeviceGridState(mIdp).getDbFile());
         try {
             // This is the current grid we have, given by the mContext
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
             // This is the state we want to migrate to that is given by the idp
-            DeviceGridState destDeviceState = new DeviceGridState(idp);
+            DeviceGridState destDeviceState = new DeviceGridState(mIdp);
 
             boolean isDestNewDb = !existingDBs.contains(destDeviceState.getDbFile());
             GridSizeMigrationLogic gridSizeMigrationLogic = new GridSizeMigrationLogic();
@@ -404,10 +410,10 @@ public class ModelDbController {
             ModelDelegate modelDelegate) {
         if (!migrateGridIfNeeded(modelDelegate)) {
             if (restoreEventLogger != null) {
-                if (LauncherPrefs.get(mContext).get(NO_DB_FILES_RESTORED)) {
+                if (mPrefs.get(NO_DB_FILES_RESTORED)) {
                     restoreEventLogger.logLauncherItemsRestoreFailed(DATA_TYPE_DB_FILE, 1,
                             RestoreError.DATABASE_FILE_NOT_RESTORED);
-                    LauncherPrefs.get(mContext).put(NO_DB_FILES_RESTORED, false);
+                    mPrefs.put(NO_DB_FILES_RESTORED, false);
                     FileLog.d(TAG, "There is no data to migrate: resetting launcher database");
                 } else {
                     restoreEventLogger.logLauncherItemsRestored(DATA_TYPE_DB_FILE, 1);
@@ -416,11 +422,10 @@ public class ModelDbController {
             }
             FileLog.d(TAG, "tryMigrateDB: Migration failed: resetting launcher database");
             createEmptyDB();
-            LauncherPrefs.get(mContext).putSync(
-                    getEmptyDbCreatedKey(mOpenHelper.getDatabaseName()).to(true));
+            mPrefs.putSync(getEmptyDbCreatedKey(mOpenHelper.getDatabaseName()).to(true));
 
             // Write the grid state to avoid another migration
-            new DeviceGridState(LauncherAppState.getIDP(mContext)).writeToPrefs(mContext);
+            new DeviceGridState(mIdp).writeToPrefs(mContext);
         } else if (restoreEventLogger != null) {
             restoreEventLogger.logLauncherItemsRestored(DATA_TYPE_DB_FILE, 1);
         }
@@ -434,17 +439,16 @@ public class ModelDbController {
      */
     private boolean migrateGridIfNeeded(ModelDelegate modelDelegate) {
         createDbIfNotExists();
-        if (LauncherPrefs.get(mContext).get(getEmptyDbCreatedKey())) {
+        if (mPrefs.get(getEmptyDbCreatedKey())) {
             // If we have already create a new DB, ignore migration
             FileLog.d(TAG, "migrateGridIfNeeded: new DB already created, skipping migration");
             return false;
         }
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
-        if (!GridSizeMigrationDBController.needsToMigrate(mContext, idp)) {
+        if (!GridSizeMigrationDBController.needsToMigrate(mContext, mIdp)) {
             FileLog.d(TAG, "migrateGridIfNeeded: no grid migration needed");
             return true;
         }
-        String targetDbName = new DeviceGridState(idp).getDbFile();
+        String targetDbName = new DeviceGridState(mIdp).getDbFile();
         if (TextUtils.equals(targetDbName, mOpenHelper.getDatabaseName())) {
             FileLog.e(TAG, "migrateGridIfNeeded: target db is same as current"
                     + " current db: " + mOpenHelper.getDatabaseName()
@@ -462,7 +466,7 @@ public class ModelDbController {
             // This is the current grid we have, given by the mContext
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
             // This is the state we want to migrate to that is given by the idp
-            DeviceGridState destDeviceState = new DeviceGridState(idp);
+            DeviceGridState destDeviceState = new DeviceGridState(mIdp);
             boolean isDestNewDb = !existingDBs.contains(destDeviceState.getDbFile());
             return GridSizeMigrationDBController.migrateGridIfNeeded(mContext, srcDeviceState,
                     destDeviceState, mOpenHelper, oldHelper.getWritableDatabase(), isDestNewDb,
@@ -611,7 +615,7 @@ public class ModelDbController {
     }
 
     private void clearFlagEmptyDbCreated() {
-        LauncherPrefs.get(mContext).removeSync(getEmptyDbCreatedKey());
+        mPrefs.removeSync(getEmptyDbCreatedKey());
     }
 
     /**
@@ -625,7 +629,7 @@ public class ModelDbController {
     public synchronized void loadDefaultFavoritesIfNecessary() {
         createDbIfNotExists();
 
-        if (LauncherPrefs.get(mContext).get(getEmptyDbCreatedKey())) {
+        if (mPrefs.get(getEmptyDbCreatedKey())) {
             Log.d(TAG, "loading default workspace");
 
             LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
@@ -737,10 +741,9 @@ public class ModelDbController {
     }
 
     private DefaultLayoutParser getDefaultLayoutParser(LauncherWidgetHolder widgetHolder) {
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
-        int defaultLayout = idp.demoModeLayoutId != 0
+        int defaultLayout = mIdp.demoModeLayoutId != 0
                 && mContext.getSystemService(UserManager.class).isDemoUser()
-                ? idp.demoModeLayoutId : idp.defaultLayoutId;
+                ? mIdp.demoModeLayoutId : mIdp.defaultLayoutId;
 
         return new DefaultLayoutParser(mContext, widgetHolder,
                 mOpenHelper, mContext.getResources(), defaultLayout);
@@ -766,6 +769,6 @@ public class ModelDbController {
      * Returns the serial number for the provided user
      */
     public long getSerialNumberForUser(UserHandle user) {
-        return UserCache.INSTANCE.get(mContext).getSerialNumberForUser(user);
+        return mUserCache.getSerialNumberForUser(user);
     }
 }
diff --git a/src/com/android/launcher3/model/ModelWriter.java b/src/com/android/launcher3/model/ModelWriter.java
index 0332775224..659bacd784 100644
--- a/src/com/android/launcher3/model/ModelWriter.java
+++ b/src/com/android/launcher3/model/ModelWriter.java
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.model;
 
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.provider.LauncherDbUtils.itemIdMatch;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
@@ -253,7 +252,7 @@ public class ModelWriter {
             item.onAddToDatabase(writer);
             writer.put(Favorites._ID, item.id);
 
-            mModel.getModelDbController().insert(Favorites.TABLE_NAME, writer.getValues(mContext));
+            mModel.getModelDbController().insert(writer.getValues(mContext));
             synchronized (mBgDataModel) {
                 checkItemInfoLocked(item.id, item, stackTrace);
                 mBgDataModel.addItem(mContext, item, true);
@@ -292,7 +291,7 @@ public class ModelWriter {
         notifyDelete(items);
         enqueueDeleteRunnable(newModelTask(() -> {
             for (ItemInfo item : items) {
-                mModel.getModelDbController().delete(TABLE_NAME, itemIdMatch(item.id), null);
+                mModel.getModelDbController().delete(itemIdMatch(item.id), null);
                 mBgDataModel.removeItem(mContext, item);
                 verifier.verifyModel();
             }
@@ -307,12 +306,12 @@ public class ModelWriter {
         notifyDelete(Collections.singleton(info));
 
         enqueueDeleteRunnable(newModelTask(() -> {
-            mModel.getModelDbController().delete(Favorites.TABLE_NAME,
+            mModel.getModelDbController().delete(
                     Favorites.CONTAINER + "=" + info.id, null);
             mBgDataModel.removeItem(mContext, info.getContents());
             info.getContents().clear();
 
-            mModel.getModelDbController().delete(Favorites.TABLE_NAME,
+            mModel.getModelDbController().delete(
                     Favorites._ID + "=" + info.id, null);
             mBgDataModel.removeItem(mContext, info);
             verifier.verifyModel();
@@ -411,7 +410,7 @@ public class ModelWriter {
         @Override
         public void runImpl() {
             mModel.getModelDbController().update(
-                    TABLE_NAME, mWriter.get().getValues(mContext), itemIdMatch(mItemId), null);
+                    mWriter.get().getValues(mContext), itemIdMatch(mItemId), null);
             updateItemArrays(mItem, mItemId);
         }
     }
@@ -433,7 +432,7 @@ public class ModelWriter {
                     ItemInfo item = mItems.get(i);
                     final int itemId = item.id;
                     mModel.getModelDbController().update(
-                            TABLE_NAME, mValues.get(i), itemIdMatch(itemId), null);
+                            mValues.get(i), itemIdMatch(itemId), null);
                     updateItemArrays(item, itemId);
                 }
                 t.commit();
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 23941bb1e6..f6ee26baa6 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -80,6 +80,7 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 /**
@@ -206,7 +207,8 @@ public class RestoreDbTask {
             LauncherRestoreEventLogger restoreEventLogger =
                     LauncherRestoreEventLogger.Companion.newInstance(context);
             task.sanitizeDB(context, controller, db, backupManager, restoreEventLogger);
-            task.restoreAppWidgetIdsIfExists(context, controller, restoreEventLogger);
+            task.restoreAppWidgetIdsIfExists(context, controller, restoreEventLogger,
+                    () -> new AppWidgetHost(context, APPWIDGET_HOST_ID));
             t.commit();
             return true;
         } catch (Exception e) {
@@ -438,14 +440,13 @@ public class RestoreDbTask {
     @WorkerThread
     @VisibleForTesting
     void restoreAppWidgetIdsIfExists(Context context, ModelDbController controller,
-            LauncherRestoreEventLogger restoreEventLogger) {
+            LauncherRestoreEventLogger restoreEventLogger, Supplier<AppWidgetHost> hostSupplier) {
         LauncherPrefs lp = LauncherPrefs.get(context);
         if (lp.has(APP_WIDGET_IDS, OLD_APP_WIDGET_IDS)) {
-            AppWidgetHost host = new AppWidgetHost(context, APPWIDGET_HOST_ID);
             restoreAppWidgetIds(context, controller, restoreEventLogger,
                     IntArray.fromConcatString(lp.get(OLD_APP_WIDGET_IDS)).toArray(),
                     IntArray.fromConcatString(lp.get(APP_WIDGET_IDS)).toArray(),
-                    host);
+                    hostSupplier.get());
         } else {
             FileLog.d(TAG, "Did not receive new app widget id map during Launcher restore");
         }
diff --git a/src/com/android/launcher3/util/ContentWriter.java b/src/com/android/launcher3/util/ContentWriter.java
index 9910dc2e70..c2c1fee561 100644
--- a/src/com/android/launcher3/util/ContentWriter.java
+++ b/src/com/android/launcher3/util/ContentWriter.java
@@ -23,7 +23,6 @@ import android.os.UserHandle;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.GraphicsUtils;
 import com.android.launcher3.model.ModelDbController;
@@ -107,7 +106,7 @@ public class ContentWriter {
     public int commit() {
         if (mCommitParams != null) {
             return mCommitParams.mDbController.update(
-                    Favorites.TABLE_NAME, getValues(mContext),
+                    getValues(mContext),
                     mCommitParams.mWhere, mCommitParams.mSelectionArgs);
         }
         return 0;
diff --git a/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java b/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
index 47110fa0e6..f7723392d4 100644
--- a/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
+++ b/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.celllayout;
 
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.TestUtil.runOnExecutorSync;
@@ -89,7 +88,7 @@ public class FavoriteItemsTransaction {
 
                     item.onAddToDatabase(writer);
                     writer.put(LauncherSettings.Favorites._ID, i);
-                    controller.insert(TABLE_NAME, writer.getValues(mContext));
+                    controller.insert(writer.getValues(mContext));
                 }
 
                 for (int i = 0; i < containerItems.size(); i++) {
@@ -97,7 +96,7 @@ public class FavoriteItemsTransaction {
                     ItemInfo item = containerItems.get(i);
                     item.onAddToDatabase(writer);
                     writer.put(LauncherSettings.Favorites._ID, count + i);
-                    controller.insert(TABLE_NAME, writer.getValues(mContext));
+                    controller.insert(writer.getValues(mContext));
                 }
                 transaction.commit();
             }
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/DatabaseHelperTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/DatabaseHelperTest.kt
index 09752b860e..89e676f52f 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/DatabaseHelperTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/DatabaseHelperTest.kt
@@ -11,6 +11,7 @@ import com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE
 import com.android.launcher3.LauncherSettings.Favorites.addTableToDb
 import com.android.launcher3.pm.UserCache
 import com.android.launcher3.provider.LauncherDbUtils
+import com.android.launcher3.util.ModelTestExtensions
 import java.util.function.ToLongFunction
 import org.junit.After
 import org.junit.Assert.assertEquals
@@ -33,7 +34,7 @@ class DatabaseHelperTest {
 
     @Before
     fun setUp() {
-        db = FactitiousDbController(context, INSERTION_SQL).inMemoryDb
+        db = ModelTestExtensions.createInMemoryDb(INSERTION_SQL)
     }
 
     @After
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FactitiousDbController.kt b/tests/multivalentTests/src/com/android/launcher3/model/FactitiousDbController.kt
deleted file mode 100644
index 711e1d2d65..0000000000
--- a/tests/multivalentTests/src/com/android/launcher3/model/FactitiousDbController.kt
+++ /dev/null
@@ -1,60 +0,0 @@
-package com.android.launcher3.model
-
-import android.content.Context
-import android.database.Cursor
-import android.database.sqlite.SQLiteDatabase
-import androidx.test.platform.app.InstrumentationRegistry
-import java.io.BufferedReader
-import java.io.InputStreamReader
-
-private val All_COLUMNS =
-    arrayOf(
-        "_id",
-        "title",
-        "intent",
-        "container",
-        "screen",
-        "cellX",
-        "cellY",
-        "spanX",
-        "spanY",
-        "itemType",
-        "appWidgetId",
-        "icon",
-        "appWidgetProvider",
-        "modified",
-        "restored",
-        "profileId",
-        "rank",
-        "options",
-        "appWidgetSource"
-    )
-
-class FactitiousDbController(context: Context, insertFile: String) : ModelDbController(context) {
-
-    val inMemoryDb: SQLiteDatabase by lazy {
-        SQLiteDatabase.createInMemory(SQLiteDatabase.OpenParams.Builder().build()).also { db ->
-            BufferedReader(
-                    InputStreamReader(
-                        InstrumentationRegistry.getInstrumentation().context.assets.open(insertFile)
-                    )
-                )
-                .lines()
-                .forEach { sqlStatement -> db.execSQL(sqlStatement) }
-        }
-    }
-
-    override fun query(
-        table: String,
-        projection: Array<out String>?,
-        selection: String?,
-        selectionArgs: Array<out String>?,
-        sortOrder: String?
-    ): Cursor {
-        return inMemoryDb.query(table, All_COLUMNS, selection, selectionArgs, null, null, sortOrder)
-    }
-
-    override fun loadDefaultFavoritesIfNecessary() {
-        // No-Op
-    }
-}
diff --git a/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index c30b730807..0f4940e7ed 100644
--- a/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -24,8 +24,6 @@ import static com.android.launcher3.LauncherPrefs.OLD_APP_WIDGET_IDS;
 import static com.android.launcher3.LauncherPrefs.RESTORE_DEVICE;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
-import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -50,7 +48,6 @@ import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.os.UserHandle;
 import android.os.UserManager;
-import android.util.LongSparseArray;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -60,9 +57,17 @@ import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.backuprestore.LauncherRestoreEventLogger;
+import com.android.launcher3.dagger.LauncherAppComponent;
+import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.model.ModelDbController;
+import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.util.AllModulesForTest;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext;
+
+import dagger.BindsInstance;
+import dagger.Component;
 
 import org.junit.After;
 import org.junit.Before;
@@ -85,7 +90,9 @@ public class RestoreDbTaskTest {
     private final UserHandle mWorkUser = UserHandle.getUserHandleForUid(PER_USER_RANGE);
 
     private LauncherModelHelper mModelHelper;
-    private Context mContext;
+    private SandboxModelContext mContext;
+    private UserCache mUserCacheSpy;
+
     private RestoreDbTask mTask;
     private ModelDbController mMockController;
     private SQLiteDatabase mMockDb;
@@ -94,10 +101,17 @@ public class RestoreDbTaskTest {
     private LauncherRestoreEventLogger mMockRestoreEventLogger;
     private SQLiteDatabase mDb;
 
+    private AppWidgetHost mWidgetHost;
+
     @Before
     public void setup() {
         mModelHelper = new LauncherModelHelper();
         mContext = mModelHelper.sandboxContext;
+        mUserCacheSpy = spy(UserCache.getInstance(getInstrumentation().getTargetContext()));
+
+        mContext.initDaggerComponent(
+                DaggerRestoreDbTaskTest_TestComponent.builder().bindUserCache(mUserCacheSpy));
+
         mTask = new RestoreDbTask();
         mMockController = Mockito.mock(ModelDbController.class);
         mMockDb = mock(SQLiteDatabase.class);
@@ -106,24 +120,34 @@ public class RestoreDbTaskTest {
         mMockRestoreEventLogger = mock(LauncherRestoreEventLogger.class);
     }
 
+    private synchronized AppWidgetHost getWidgetHostLazy() {
+        if (mWidgetHost == null) {
+            mWidgetHost = new AppWidgetHost(mContext, 1012);
+        }
+        return mWidgetHost;
+    }
+
     @After
     public void teardown() {
         if (mDb != null) {
             mDb.close();
         }
+        if (mWidgetHost != null) {
+            mWidgetHost.deleteHost();
+        }
         mModelHelper.destroy();
         LauncherPrefs.get(mContext).removeSync(RESTORE_DEVICE);
     }
 
     @Test
     public void testGetProfileId() throws Exception {
-        mDb = new MyModelDbController(23).getDb();
+        mDb = getModelDbController(23).getDb();
         assertEquals(23, new RestoreDbTask().getDefaultProfileId(mDb));
     }
 
     @Test
     public void testMigrateProfileId() throws Exception {
-        mDb = new MyModelDbController(42).getDb();
+        mDb = getModelDbController(42).getDb();
         // Add some mock data
         for (int i = 0; i < 5; i++) {
             ContentValues values = new ContentValues();
@@ -143,7 +167,7 @@ public class RestoreDbTaskTest {
 
     @Test
     public void testChangeDefaultColumn() throws Exception {
-        mDb = new MyModelDbController(42).getDb();
+        mDb = getModelDbController(42).getDb();
         // Add some mock data
         for (int i = 0; i < 5; i++) {
             ContentValues values = new ContentValues();
@@ -173,12 +197,12 @@ public class RestoreDbTaskTest {
         long workProfileId = myProfileId + 2;
         long workProfileId_old = myProfileId + 3;
 
-        MyModelDbController controller = new MyModelDbController(myProfileId);
+        ModelDbController controller = getModelDbController(myProfileId);
         mDb = controller.getDb();
         BackupManager bm = spy(new BackupManager(mContext));
         doReturn(myUserHandle()).when(bm).getUserForAncestralSerialNumber(eq(myProfileId_old));
         doReturn(mWorkUser).when(bm).getUserForAncestralSerialNumber(eq(workProfileId_old));
-        controller.users.put(workProfileId, mWorkUser);
+        doReturn(workProfileId).when(mUserCacheSpy).getSerialNumberForUser(mWorkUser);
 
         addIconsBulk(controller, 10, 1, myProfileId_old);
         addIconsBulk(controller, 6, 2, workProfileId_old);
@@ -202,7 +226,7 @@ public class RestoreDbTaskTest {
         long myProfileId_old = myProfileId + 1;
         long workProfileId_old = myProfileId + 3;
 
-        MyModelDbController controller = new MyModelDbController(myProfileId);
+        ModelDbController controller = getModelDbController(myProfileId);
         mDb = controller.getDb();
         BackupManager bm = spy(new BackupManager(mContext));
         doReturn(myUserHandle()).when(bm).getUserForAncestralSerialNumber(eq(myProfileId_old));
@@ -226,7 +250,8 @@ public class RestoreDbTaskTest {
     @Test
     public void givenLauncherPrefsHasNoIds_whenRestoreAppWidgetIdsIfExists_thenIdsAreRemoved() {
         // When
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger,
+                this::getWidgetHostLazy);
         // Then
         assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
     }
@@ -234,7 +259,7 @@ public class RestoreDbTaskTest {
     @Test
     public void givenNoPendingRestore_WhenRestoreAppWidgetIds_ThenRemoveNewWidgetIds() {
         // Given
-        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
+        AppWidgetHost expectedHost = getWidgetHostLazy();
         int[] expectedOldIds = generateOldWidgetIds(expectedHost);
         int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
         when(mMockController.getDb()).thenReturn(mMockDb);
@@ -242,7 +267,8 @@ public class RestoreDbTaskTest {
 
         // When
         setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger,
+                this::getWidgetHostLazy);
 
         // Then
         assertThat(expectedHost.getAppWidgetIds()).isEqualTo(expectedOldIds);
@@ -254,7 +280,7 @@ public class RestoreDbTaskTest {
     @Test
     public void givenRestoreWithNonExistingWidgets_WhenRestoreAppWidgetIds_ThenRemoveNewIds() {
         // Given
-        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
+        AppWidgetHost expectedHost = getWidgetHostLazy();
         int[] expectedOldIds = generateOldWidgetIds(expectedHost);
         int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
         when(mMockController.getDb()).thenReturn(mMockDb);
@@ -265,18 +291,19 @@ public class RestoreDbTaskTest {
 
         // When
         setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger,
+                this::getWidgetHostLazy);
 
         // Then
         assertThat(expectedHost.getAppWidgetIds()).isEqualTo(expectedOldIds);
         assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
-        verify(mMockController, times(expectedOldIds.length)).update(any(), any(), any(), any());
+        verify(mMockController, times(expectedOldIds.length)).update(any(), any(), any());
     }
 
     @Test
     public void givenRestore_WhenRestoreAppWidgetIds_ThenAddNewIds() {
         // Given
-        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
+        AppWidgetHost expectedHost = getWidgetHostLazy();
         int[] expectedOldIds = generateOldWidgetIds(expectedHost);
         int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
         int[] allExpectedIds = IntStream.concat(
@@ -294,15 +321,16 @@ public class RestoreDbTaskTest {
 
         // When
         setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController, mMockRestoreEventLogger,
+                this::getWidgetHostLazy);
 
         // Then
         assertThat(expectedHost.getAppWidgetIds()).isEqualTo(allExpectedIds);
         assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
-        verify(mMockController, times(expectedOldIds.length)).update(any(), any(), any(), any());
+        verify(mMockController, times(expectedOldIds.length)).update(any(), any(), any());
     }
 
-    private void addIconsBulk(MyModelDbController controller,
+    private void addIconsBulk(ModelDbController controller,
             int count, int screen, long profileId) {
         int columns = LauncherAppState.getIDP(mContext).numColumns;
         String packageName = getInstrumentation().getContext().getPackageName();
@@ -320,7 +348,7 @@ public class RestoreDbTaskTest {
             values.put(LauncherSettings.Favorites.INTENT,
                     new Intent(Intent.ACTION_MAIN).setPackage(packageName).toUri(0));
 
-            controller.insert(TABLE_NAME, values);
+            controller.insert(values);
         }
     }
 
@@ -346,7 +374,7 @@ public class RestoreDbTaskTest {
     }
 
     private void runRemoveScreenIdGapsTest(int[] screenIds, int[] expectedScreenIds) {
-        mDb = new MyModelDbController(42).getDb();
+        mDb = getModelDbController(42).getDb();
         // Add some mock data
         for (int i = 0; i < screenIds.length; i++) {
             ContentValues values = new ContentValues();
@@ -397,25 +425,29 @@ public class RestoreDbTaskTest {
                 .map(id -> host.allocateAppWidgetId()).toArray();
     }
 
-    private class MyModelDbController extends ModelDbController {
-
-        public final LongSparseArray<UserHandle> users = new LongSparseArray<>();
-
-        MyModelDbController(long profileId) {
-            super(mContext);
-            users.put(profileId, myUserHandle());
-        }
-
-        @Override
-        public long getSerialNumberForUser(UserHandle user) {
-            int index = users.indexOfValue(user);
-            return index >= 0 ? users.keyAt(index) : -1;
-        }
-    }
-
     private void setRestoredAppWidgetIds(Context context, int[] oldIds, int[] newIds) {
         LauncherPrefs.get(context).putSync(
                 OLD_APP_WIDGET_IDS.to(IntArray.wrap(oldIds).toConcatString()),
                 APP_WIDGET_IDS.to(IntArray.wrap(newIds).toConcatString()));
     }
+
+    private ModelDbController getModelDbController(long profileId) {
+        doReturn(profileId).when(mUserCacheSpy).getSerialNumberForUser(myUserHandle());
+        return ((TestComponent) mContext.getAppComponent()).getDbController();
+    }
+
+    @LauncherAppSingleton
+    @Component(modules = AllModulesForTest.class)
+    public interface TestComponent extends LauncherAppComponent {
+
+        ModelDbController getDbController();
+
+        @Component.Builder
+        interface Builder extends LauncherAppComponent.Builder {
+
+            @BindsInstance Builder bindUserCache(UserCache userCache);
+
+            TestComponent build();
+        }
+    }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt b/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
index ceefb0dced..524acff61d 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
@@ -1,10 +1,11 @@
 package com.android.launcher3.util
 
 import android.content.ContentValues
+import android.database.sqlite.SQLiteDatabase
 import android.os.Process
+import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.Flags
 import com.android.launcher3.LauncherModel
-import com.android.launcher3.LauncherSettings.Favorites
 import com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_ID
 import com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_PROVIDER
 import com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_SOURCE
@@ -24,6 +25,8 @@ import com.android.launcher3.LauncherSettings.Favorites.TITLE
 import com.android.launcher3.LauncherSettings.Favorites._ID
 import com.android.launcher3.model.BgDataModel
 import com.android.launcher3.model.ModelDbController
+import java.io.BufferedReader
+import java.io.InputStreamReader
 
 object ModelTestExtensions {
     /** Clears and reloads Launcher db to cleanup the workspace */
@@ -68,7 +71,6 @@ object ModelTestExtensions {
         spanY: Int = 1,
         id: Int = 0,
         profileId: Int = Process.myUserHandle().identifier,
-        tableName: String = Favorites.TABLE_NAME,
         appWidgetId: Int = -1,
         appWidgetSource: Int = -1,
         appWidgetProvider: String? = null,
@@ -97,9 +99,21 @@ object ModelTestExtensions {
                         values[APPWIDGET_PROVIDER] = appWidgetProvider
                     }
                 // Migrate any previous data so that the DB state is correct
-                controller.insert(tableName, values)
+                controller.insert(values)
                 transaction.commit()
             }
         }
     }
+
+    /** Creates an in-memory sqlite DB and initializes with the data in [insertFile] */
+    fun createInMemoryDb(insertFile: String): SQLiteDatabase =
+        SQLiteDatabase.createInMemory(SQLiteDatabase.OpenParams.Builder().build()).also { db ->
+            BufferedReader(
+                    InputStreamReader(
+                        InstrumentationRegistry.getInstrumentation().context.assets.open(insertFile)
+                    )
+                )
+                .lines()
+                .forEach { sqlStatement -> db.execSQL(sqlStatement) }
+        }
 }
diff --git a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
index 38fad6bc6c..34d9d40d67 100644
--- a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
+++ b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
@@ -23,8 +23,8 @@ import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.Flags
+import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherPrefs
-import com.android.launcher3.model.ModelDbController
 import com.android.launcher3.model.ModelDelegate
 import com.android.launcher3.provider.RestoreDbTask
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
@@ -45,11 +45,8 @@ import org.mockito.kotlin.mock
 @MediumTest
 class BackupAndRestoreDBSelectionTest {
 
-    @JvmField @Rule var backAndRestoreRule = BackAndRestoreRule()
-
-    @JvmField
-    @Rule
-    val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+    @get:Rule var backAndRestoreRule = BackAndRestoreRule()
+    @get:Rule val setFlagsRule = SetFlagsRule()
 
     val modelDelegate = mock<ModelDelegate>()
 
@@ -70,7 +67,8 @@ class BackupAndRestoreDBSelectionTest {
 
     @Test
     fun oldDatabasesNotPresentAfterRestore() {
-        val dbController = ModelDbController(getInstrumentation().targetContext)
+        val dbController =
+            LauncherAppState.getInstance(getInstrumentation().targetContext).model.modelDbController
         if (Flags.gridMigrationRefactor()) {
             dbController.attemptMigrateDb(null, modelDelegate)
         } else {
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index 7f9b7a0ecb..560d306244 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -5,6 +5,7 @@ import android.content.ComponentName
 import android.content.Intent
 import android.content.pm.ApplicationInfo
 import android.content.pm.LauncherActivityInfo
+import android.database.sqlite.SQLiteDatabase
 import android.os.Process
 import android.os.UserHandle
 import android.platform.test.annotations.DisableFlags
@@ -25,6 +26,7 @@ import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
 import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR
 import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
+import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
 import com.android.launcher3.dagger.LauncherAppComponent
 import com.android.launcher3.dagger.LauncherAppSingleton
 import com.android.launcher3.icons.IconCache
@@ -41,6 +43,7 @@ import com.android.launcher3.util.AllModulesForTest
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
 import com.android.launcher3.util.LooperIdleLock
+import com.android.launcher3.util.ModelTestExtensions
 import com.android.launcher3.util.TestUtil
 import com.android.launcher3.util.UserIconInfo
 import com.google.common.truth.Truth.assertThat
@@ -63,6 +66,7 @@ import org.mockito.MockitoSession
 import org.mockito.Spy
 import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
+import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
@@ -94,6 +98,7 @@ class LoaderTaskTest {
     @Mock private lateinit var launcherModel: LauncherModel
     @Mock private lateinit var iconCache: IconCache
     @Mock private lateinit var userCache: UserCache
+    @Mock private lateinit var modelDbController: ModelDbController
 
     @Mock private lateinit var launcherBinder: BaseLauncherBinder
     @Mock private lateinit var transaction: LoaderTransaction
@@ -110,6 +115,10 @@ class LoaderTaskTest {
     private val bgDataModel: BgDataModel
         get() = testComponent.getDataModel()
 
+    private val inMemoryDb: SQLiteDatabase by lazy {
+        ModelTestExtensions.createInMemoryDb(INSERTION_STATEMENT_FILE)
+    }
+
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
@@ -123,8 +132,23 @@ class LoaderTaskTest {
             .getAppWidgetInfo(any())
 
         `when`(launcherModel.beginLoader(any())).thenReturn(transaction)
-        `when`(launcherModel.modelDbController)
-            .thenReturn(FactitiousDbController(context, INSERTION_STATEMENT_FILE))
+
+        `when`(launcherModel.modelDbController).thenReturn(modelDbController)
+        doAnswer {}.whenever(modelDbController).loadDefaultFavoritesIfNecessary()
+        doAnswer { i ->
+                inMemoryDb.query(
+                    TABLE_NAME,
+                    i.getArgument(0),
+                    i.getArgument(1),
+                    i.getArgument(2),
+                    null,
+                    null,
+                    i.getArgument(3),
+                )
+            }
+            .whenever(modelDbController)
+            .query(anyOrNull(), anyOrNull(), anyOrNull(), anyOrNull())
+
         `when`(launcherModel.modelDelegate).thenReturn(modelDelegate)
         `when`(launcherBinder.newIdleLock(any())).thenReturn(idleLock)
         `when`(idleLock.awaitLocked(1000)).thenReturn(false)
@@ -149,6 +173,7 @@ class LoaderTaskTest {
     fun tearDown() {
         LauncherPrefs.get(context).removeSync(RESTORE_DEVICE)
         LauncherPrefs.get(context).putSync(IS_FIRST_LOAD_AFTER_RESTORE.to(false))
+        inMemoryDb.close()
         context.onDestroy()
         mockitoSession.finishMocking()
     }
diff --git a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
index 8846d6593a..67ee1ac789 100644
--- a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.ui.widget;
 
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
 import static com.android.launcher3.model.data.LauncherAppWidgetInfo.FLAG_ID_NOT_VALID;
 import static com.android.launcher3.model.data.LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
@@ -318,7 +317,7 @@ public class BindWidgetTest extends BaseLauncherActivityTest<Launcher> {
         try {
             return MODEL_EXECUTOR.submit(() ->
                 mModel.getModelDbController().query(
-                                TABLE_NAME, null, itemIdMatch(0), null, null)).get();
+                        null, itemIdMatch(0), null, null)).get();
         } catch (Exception e) {
             throw new RuntimeException(e);
         }