Move AppWidgetsRestoredReceiver methods to enable adding tests for restoring widget Ids

Bug: 294386159
Test: Ran unit tests locally
Change-Id: Ib9657e6a0faa876197e0106b8729dcc606562d09
diff --git a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
index ec874b9..55b8fcc 100644
--- a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
+++ b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
@@ -1,28 +1,12 @@
 package com.android.launcher3;
 
-import static android.os.Process.myUserHandle;
-
-import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetManager;
-import android.appwidget.AppWidgetProviderInfo;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
-import android.database.Cursor;
 import android.util.Log;
 
-import androidx.annotation.NonNull;
-import androidx.annotation.WorkerThread;
-
-import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.model.LoaderTask;
-import com.android.launcher3.model.ModelDbController;
-import com.android.launcher3.model.WidgetsModel;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.RestoreDbTask;
-import com.android.launcher3.util.ContentWriter;
-import com.android.launcher3.util.IntArray;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
 public class AppWidgetsRestoredReceiver extends BroadcastReceiver {
@@ -47,131 +31,4 @@
             }
         }
     }
-
-    /**
-     * Updates the app widgets whose id has changed during the restore process.
-     */
-    @WorkerThread
-    public static void restoreAppWidgetIds(Context context, ModelDbController controller,
-            int[] oldWidgetIds, int[] newWidgetIds, @NonNull AppWidgetHost host) {
-        if (WidgetsModel.GO_DISABLE_WIDGETS) {
-            Log.e(TAG, "Skipping widget ID remap as widgets not supported");
-            host.deleteHost();
-            return;
-        }
-        if (!RestoreDbTask.isPending(context)) {
-            // Someone has already gone through our DB once, probably LoaderTask. Skip any further
-            // modifications of the DB.
-            Log.e(TAG, "Skipping widget ID remap as DB already in use");
-            for (int widgetId : newWidgetIds) {
-                Log.d(TAG, "Deleting widgetId: " + widgetId);
-                host.deleteAppWidgetId(widgetId);
-            }
-            return;
-        }
-
-        final AppWidgetManager widgets = AppWidgetManager.getInstance(context);
-
-        Log.d(TAG, "restoreAppWidgetIds: "
-                + "oldWidgetIds=" + IntArray.wrap(oldWidgetIds).toConcatString()
-                + ", newWidgetIds=" + IntArray.wrap(newWidgetIds).toConcatString());
-
-        // TODO(b/234700507): Remove the logs after the bug is fixed
-        logDatabaseWidgetInfo(controller);
-
-        for (int i = 0; i < oldWidgetIds.length; i++) {
-            Log.i(TAG, "Widget state restore id " + oldWidgetIds[i] + " => " + newWidgetIds[i]);
-
-            final AppWidgetProviderInfo provider = widgets.getAppWidgetInfo(newWidgetIds[i]);
-            final int state;
-            if (LoaderTask.isValidProvider(provider)) {
-                // This will ensure that we show 'Click to setup' UI if required.
-                state = LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
-            } else {
-                state = LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
-            }
-
-            // b/135926478: Work profile widget restore is broken in platform. This forces us to
-            // recreate the widget during loading with the correct host provider.
-            long mainProfileId = UserCache.INSTANCE.get(context)
-                    .getSerialNumberForUser(myUserHandle());
-            long controllerProfileId = controller.getSerialNumberForUser(myUserHandle());
-            String oldWidgetId = Integer.toString(oldWidgetIds[i]);
-            final String where = "appWidgetId=? and (restored & 1) = 1 and profileId=?";
-            String profileId = Long.toString(mainProfileId);
-            final String[] args = new String[] { oldWidgetId, profileId };
-            Log.d(TAG, "restoreAppWidgetIds: querying profile id=" + profileId
-                    + " with controller profile ID=" + controllerProfileId);
-            int result = new ContentWriter(context,
-                            new ContentWriter.CommitParams(controller, where, args))
-                    .put(LauncherSettings.Favorites.APPWIDGET_ID, newWidgetIds[i])
-                    .put(LauncherSettings.Favorites.RESTORED, state)
-                    .commit();
-            if (result == 0) {
-                // TODO(b/234700507): Remove the logs after the bug is fixed
-                Log.e(TAG, "restoreAppWidgetIds: remapping failed since the widget is not in"
-                        + " the database anymore");
-                try (Cursor cursor = controller.getDb().query(
-                        Favorites.TABLE_NAME,
-                        new String[]{Favorites.APPWIDGET_ID},
-                        "appWidgetId=?", new String[]{oldWidgetId}, null, null, null)) {
-                    if (!cursor.moveToFirst()) {
-                        // The widget no long exists.
-                        Log.d(TAG, "Deleting widgetId: " + newWidgetIds[i] + " with old id: "
-                                + oldWidgetId);
-                        host.deleteAppWidgetId(newWidgetIds[i]);
-                    }
-                }
-            }
-        }
-
-        LauncherAppState app = LauncherAppState.getInstanceNoCreate();
-        if (app != null) {
-            app.getModel().forceReload();
-        }
-    }
-
-    private static void logDatabaseWidgetInfo(ModelDbController controller) {
-        try (Cursor cursor = controller.getDb().query(Favorites.TABLE_NAME,
-                new String[]{Favorites.APPWIDGET_ID, Favorites.RESTORED, Favorites.PROFILE_ID},
-                Favorites.APPWIDGET_ID + "!=" + LauncherAppWidgetInfo.NO_ID, null,
-                null, null, null)) {
-            IntArray widgetIdList = new IntArray();
-            IntArray widgetRestoreList = new IntArray();
-            IntArray widgetProfileIdList = new IntArray();
-
-            if (cursor.moveToFirst()) {
-                final int widgetIdColumnIndex = cursor.getColumnIndex(Favorites.APPWIDGET_ID);
-                final int widgetRestoredColumnIndex = cursor.getColumnIndex(Favorites.RESTORED);
-                final int widgetProfileIdIndex = cursor.getColumnIndex(Favorites.PROFILE_ID);
-                while (!cursor.isAfterLast()) {
-                    int widgetId = cursor.getInt(widgetIdColumnIndex);
-                    int widgetRestoredFlag = cursor.getInt(widgetRestoredColumnIndex);
-                    int widgetProfileId = cursor.getInt(widgetProfileIdIndex);
-
-                    widgetIdList.add(widgetId);
-                    widgetRestoreList.add(widgetRestoredFlag);
-                    widgetProfileIdList.add(widgetProfileId);
-                    cursor.moveToNext();
-                }
-            }
-
-            StringBuilder builder = new StringBuilder();
-            builder.append("[");
-            for (int i = 0; i < widgetIdList.size(); i++) {
-                builder.append("[")
-                        .append(widgetIdList.get(i))
-                        .append(", ")
-                        .append(widgetRestoreList.get(i))
-                        .append(", ")
-                        .append(widgetProfileIdList.get(i))
-                        .append("]");
-            }
-            builder.append("]");
-            Log.d(TAG, "restoreAppWidgetIds: all widget ids in database: "
-                    + builder.toString());
-        } catch (Exception ex) {
-            Log.e(TAG, "Getting widget ids from the database failed", ex);
-        }
-    }
-}
+}
\ No newline at end of file
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 4725dd1..10005e5 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -28,6 +28,8 @@
 
 import android.app.backup.BackupManager;
 import android.appwidget.AppWidgetHost;
+import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
@@ -42,20 +44,26 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
+import androidx.annotation.WorkerThread;
 
-import com.android.launcher3.AppWidgetsRestoredReceiver;
 import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
+import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.DeviceGridState;
+import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelDbController;
+import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.uioverrides.ApiWrapper;
+import com.android.launcher3.util.ContentWriter;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LogConfig;
 
@@ -377,11 +385,13 @@
                 .putSync(RESTORE_DEVICE.to(new DeviceGridState(context).getDeviceType()));
     }
 
-    private void restoreAppWidgetIdsIfExists(Context context, ModelDbController controller) {
+    @WorkerThread
+    @VisibleForTesting
+    void restoreAppWidgetIdsIfExists(Context context, ModelDbController controller) {
         LauncherPrefs lp = LauncherPrefs.get(context);
         if (lp.has(APP_WIDGET_IDS, OLD_APP_WIDGET_IDS)) {
             AppWidgetHost host = new AppWidgetHost(context, APPWIDGET_HOST_ID);
-            AppWidgetsRestoredReceiver.restoreAppWidgetIds(context, controller,
+            restoreAppWidgetIds(context, controller,
                     IntArray.fromConcatString(lp.get(OLD_APP_WIDGET_IDS)).toArray(),
                     IntArray.fromConcatString(lp.get(APP_WIDGET_IDS)).toArray(),
                     host);
@@ -392,6 +402,133 @@
         lp.remove(APP_WIDGET_IDS, OLD_APP_WIDGET_IDS);
     }
 
+    /**
+     * Updates the app widgets whose id has changed during the restore process.
+     */
+    @WorkerThread
+    private void restoreAppWidgetIds(Context context, ModelDbController controller,
+            int[] oldWidgetIds, int[] newWidgetIds, @NonNull AppWidgetHost host) {
+        if (WidgetsModel.GO_DISABLE_WIDGETS) {
+            Log.e(TAG, "Skipping widget ID remap as widgets not supported");
+            host.deleteHost();
+            return;
+        }
+        if (!RestoreDbTask.isPending(context)) {
+            // Someone has already gone through our DB once, probably LoaderTask. Skip any further
+            // modifications of the DB.
+            Log.e(TAG, "Skipping widget ID remap as DB already in use");
+            for (int widgetId : newWidgetIds) {
+                Log.d(TAG, "Deleting widgetId: " + widgetId);
+                host.deleteAppWidgetId(widgetId);
+            }
+            return;
+        }
+
+        final AppWidgetManager widgets = AppWidgetManager.getInstance(context);
+
+        Log.d(TAG, "restoreAppWidgetIds: "
+                + "oldWidgetIds=" + IntArray.wrap(oldWidgetIds).toConcatString()
+                + ", newWidgetIds=" + IntArray.wrap(newWidgetIds).toConcatString());
+
+        // TODO(b/234700507): Remove the logs after the bug is fixed
+        logDatabaseWidgetInfo(controller);
+
+        for (int i = 0; i < oldWidgetIds.length; i++) {
+            Log.i(TAG, "Widget state restore id " + oldWidgetIds[i] + " => " + newWidgetIds[i]);
+
+            final AppWidgetProviderInfo provider = widgets.getAppWidgetInfo(newWidgetIds[i]);
+            final int state;
+            if (LoaderTask.isValidProvider(provider)) {
+                // This will ensure that we show 'Click to setup' UI if required.
+                state = LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
+            } else {
+                state = LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
+            }
+
+            // b/135926478: Work profile widget restore is broken in platform. This forces us to
+            // recreate the widget during loading with the correct host provider.
+            long mainProfileId = UserCache.INSTANCE.get(context)
+                    .getSerialNumberForUser(myUserHandle());
+            long controllerProfileId = controller.getSerialNumberForUser(myUserHandle());
+            String oldWidgetId = Integer.toString(oldWidgetIds[i]);
+            final String where = "appWidgetId=? and (restored & 1) = 1 and profileId=?";
+            String profileId = Long.toString(mainProfileId);
+            final String[] args = new String[] { oldWidgetId, profileId };
+            Log.d(TAG, "restoreAppWidgetIds: querying profile id=" + profileId
+                    + " with controller profile ID=" + controllerProfileId);
+            int result = new ContentWriter(context,
+                    new ContentWriter.CommitParams(controller, where, args))
+                    .put(LauncherSettings.Favorites.APPWIDGET_ID, newWidgetIds[i])
+                    .put(LauncherSettings.Favorites.RESTORED, state)
+                    .commit();
+            if (result == 0) {
+                // TODO(b/234700507): Remove the logs after the bug is fixed
+                Log.e(TAG, "restoreAppWidgetIds: remapping failed since the widget is not in"
+                        + " the database anymore");
+                try (Cursor cursor = controller.getDb().query(
+                        Favorites.TABLE_NAME,
+                        new String[]{Favorites.APPWIDGET_ID},
+                        "appWidgetId=?", new String[]{oldWidgetId}, null, null, null)) {
+                    if (!cursor.moveToFirst()) {
+                        // The widget no long exists.
+                        Log.d(TAG, "Deleting widgetId: " + newWidgetIds[i] + " with old id: "
+                                + oldWidgetId);
+                        host.deleteAppWidgetId(newWidgetIds[i]);
+                    }
+                }
+            }
+        }
+
+        LauncherAppState app = LauncherAppState.getInstanceNoCreate();
+        if (app != null) {
+            app.getModel().forceReload();
+        }
+    }
+
+    private static void logDatabaseWidgetInfo(ModelDbController controller) {
+        try (Cursor cursor = controller.getDb().query(Favorites.TABLE_NAME,
+                new String[]{Favorites.APPWIDGET_ID, Favorites.RESTORED, Favorites.PROFILE_ID},
+                Favorites.APPWIDGET_ID + "!=" + LauncherAppWidgetInfo.NO_ID, null,
+                null, null, null)) {
+            IntArray widgetIdList = new IntArray();
+            IntArray widgetRestoreList = new IntArray();
+            IntArray widgetProfileIdList = new IntArray();
+
+            if (cursor.moveToFirst()) {
+                final int widgetIdColumnIndex = cursor.getColumnIndex(Favorites.APPWIDGET_ID);
+                final int widgetRestoredColumnIndex = cursor.getColumnIndex(Favorites.RESTORED);
+                final int widgetProfileIdIndex = cursor.getColumnIndex(Favorites.PROFILE_ID);
+                while (!cursor.isAfterLast()) {
+                    int widgetId = cursor.getInt(widgetIdColumnIndex);
+                    int widgetRestoredFlag = cursor.getInt(widgetRestoredColumnIndex);
+                    int widgetProfileId = cursor.getInt(widgetProfileIdIndex);
+
+                    widgetIdList.add(widgetId);
+                    widgetRestoreList.add(widgetRestoredFlag);
+                    widgetProfileIdList.add(widgetProfileId);
+                    cursor.moveToNext();
+                }
+            }
+
+            StringBuilder builder = new StringBuilder();
+            builder.append("[");
+            for (int i = 0; i < widgetIdList.size(); i++) {
+                builder.append("[")
+                        .append(widgetIdList.get(i))
+                        .append(", ")
+                        .append(widgetRestoreList.get(i))
+                        .append(", ")
+                        .append(widgetProfileIdList.get(i))
+                        .append("]");
+            }
+            builder.append("]");
+            Log.d(TAG, "restoreAppWidgetIds: all widget ids in database: "
+                    + builder.toString());
+        } catch (Exception ex) {
+            Log.e(TAG, "Getting widget ids from the database failed", ex);
+        }
+    }
+
     public static void setRestoredAppWidgetIds(Context context, @NonNull int[] oldIds,
             @NonNull int[] newIds) {
         LauncherPrefs.get(context).putSync(
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 73bb586..15f6177 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -19,17 +19,30 @@
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.launcher3.LauncherPrefs.APP_WIDGET_IDS;
+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;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
 
 import android.app.backup.BackupManager;
+import android.appwidget.AppWidgetHost;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
@@ -43,6 +56,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.ModelDbController;
@@ -52,6 +66,10 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mockito;
+
+import java.util.Arrays;
+import java.util.stream.IntStream;
 
 /**
  * Tests for {@link RestoreDbTask}
@@ -66,11 +84,21 @@
 
     private LauncherModelHelper mModelHelper;
     private Context mContext;
+    private RestoreDbTask mTask;
+    private ModelDbController mMockController;
+    private SQLiteDatabase mMockDb;
+    private Cursor mMockCursor;
+    private LauncherPrefs mPrefs;
 
     @Before
     public void setup() {
         mModelHelper = new LauncherModelHelper();
         mContext = mModelHelper.sandboxContext;
+        mTask = new RestoreDbTask();
+        mMockController = Mockito.mock(ModelDbController.class);
+        mMockDb = mock(SQLiteDatabase.class);
+        mMockCursor = mock(Cursor.class);
+        mPrefs = new LauncherPrefs(mContext);
     }
 
     @After
@@ -148,8 +176,7 @@
         assertEquals(10, getItemCountForProfile(db, myProfileId_old));
         assertEquals(6, getItemCountForProfile(db, workProfileId_old));
 
-        RestoreDbTask task = new RestoreDbTask();
-        task.sanitizeDB(mContext, controller, controller.getDb(), bm);
+        mTask.sanitizeDB(mContext, controller, controller.getDb(), bm);
 
         // All the data has been migrated to the new user ids
         assertEquals(0, getItemCountForProfile(db, myProfileId_old));
@@ -178,8 +205,7 @@
         assertEquals(10, getItemCountForProfile(db, myProfileId_old));
         assertEquals(6, getItemCountForProfile(db, workProfileId_old));
 
-        RestoreDbTask task = new RestoreDbTask();
-        task.sanitizeDB(mContext, controller, controller.getDb(), bm);
+        mTask.sanitizeDB(mContext, controller, controller.getDb(), bm);
 
         // All the data has been migrated to the new user ids
         assertEquals(0, getItemCountForProfile(db, myProfileId_old));
@@ -188,6 +214,83 @@
         assertEquals(10, getCount(db, "select * from favorites"));
     }
 
+    @Test
+    public void givenLauncherPrefsHasNoIds_whenRestoreAppWidgetIdsIfExists_thenIdsAreRemoved() {
+        // When
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
+        // Then
+        assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
+    }
+
+    @Test
+    public void givenNoPendingRestore_WhenRestoreAppWidgetIds_ThenRemoveNewWidgetIds() {
+        // Given
+        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
+        int[] expectedOldIds = generateOldWidgetIds(expectedHost);
+        int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
+        when(mMockController.getDb()).thenReturn(mMockDb);
+        mPrefs.remove(RESTORE_DEVICE);
+
+        // When
+        RestoreDbTask.setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
+
+        // Then
+        assertThat(expectedHost.getAppWidgetIds()).isEqualTo(expectedOldIds);
+        assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
+        verifyZeroInteractions(mMockController);
+    }
+
+    @Test
+    public void givenRestoreWithNonExistingWidgets_WhenRestoreAppWidgetIds_ThenRemoveNewIds() {
+        // Given
+        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
+        int[] expectedOldIds = generateOldWidgetIds(expectedHost);
+        int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
+        when(mMockController.getDb()).thenReturn(mMockDb);
+        when(mMockDb.query(any(), any(), any(), any(), any(), any(), any())).thenReturn(
+                mMockCursor);
+        when(mMockCursor.moveToFirst()).thenReturn(false);
+        RestoreDbTask.setPending(mContext);
+
+        // When
+        RestoreDbTask.setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
+
+        // 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());
+    }
+
+    @Test
+    public void givenRestore_WhenRestoreAppWidgetIds_ThenAddNewIds() {
+        // Given
+        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
+        int[] expectedOldIds = generateOldWidgetIds(expectedHost);
+        int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
+        int[] allExpectedIds = IntStream.concat(
+                Arrays.stream(expectedOldIds),
+                Arrays.stream(expectedNewIds)
+        ).toArray();
+
+        when(mMockController.getDb()).thenReturn(mMockDb);
+        when(mMockDb.query(any(), any(), any(), any(), any(), any(), any()))
+                .thenReturn(mMockCursor);
+        when(mMockCursor.moveToFirst()).thenReturn(true);
+        when(mMockCursor.isAfterLast()).thenReturn(true);
+        RestoreDbTask.setPending(mContext);
+
+        // When
+        RestoreDbTask.setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
+        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
+
+        // 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());
+    }
+
     private void addIconsBulk(MyModelDbController controller,
             int count, int screen, long profileId) {
         int columns = LauncherAppState.getIDP(mContext).numColumns;
@@ -270,6 +373,19 @@
         }
     }
 
+    private int[] generateOldWidgetIds(AppWidgetHost host) {
+        // generate some widget ids in case there are none
+        host.allocateAppWidgetId();
+        host.allocateAppWidgetId();
+        return host.getAppWidgetIds();
+    }
+
+    private int[] generateNewWidgetIds(AppWidgetHost host, int[] oldWidgetIds) {
+        // map as many new ids as old ids
+        return Arrays.stream(oldWidgetIds)
+                .map(id -> host.allocateAppWidgetId()).toArray();
+    }
+
     private class MyModelDbController extends ModelDbController {
 
         public final LongSparseArray<UserHandle> users = new LongSparseArray<>();